From 226da873d05cdc2467cf5ab682d945cd86f1f47c Mon Sep 17 00:00:00 2001 From: Ryan Date: Mon, 10 Aug 2026 19:25:20 -0400 Subject: [PATCH] Story 8 --- config.yaml | 1 + src/active_learning/__init__.py | 1 + src/active_learning/label_ingestor.py | 217 ++++++++++ src/active_learning/pipeline.py | 82 ++++ src/active_learning/registry.py | 223 +++++++++++ src/active_learning/trainer.py | 344 ++++++++++++++++ src/active_learning/validator.py | 204 ++++++++++ tests/test_story_08.py | 550 ++++++++++++++++++++++++++ 8 files changed, 1622 insertions(+) create mode 100644 src/active_learning/__init__.py create mode 100644 src/active_learning/label_ingestor.py create mode 100644 src/active_learning/pipeline.py create mode 100644 src/active_learning/registry.py create mode 100644 src/active_learning/trainer.py create mode 100644 src/active_learning/validator.py create mode 100644 tests/test_story_08.py diff --git a/config.yaml b/config.yaml index 3ee6316..f8eec6d 100644 --- a/config.yaml +++ b/config.yaml @@ -184,6 +184,7 @@ review_ui: active_learning: enabled: true min_annotated_samples: 100 + seed: 42 # RNG seed for deterministic train/val splits and training training: epochs: 20 batch_size: 32 diff --git a/src/active_learning/__init__.py b/src/active_learning/__init__.py new file mode 100644 index 0000000..66758f3 --- /dev/null +++ b/src/active_learning/__init__.py @@ -0,0 +1 @@ +"""Active learning pipeline: label ingestion, training, validation, and deployment.""" diff --git a/src/active_learning/label_ingestor.py b/src/active_learning/label_ingestor.py new file mode 100644 index 0000000..12fe7e4 --- /dev/null +++ b/src/active_learning/label_ingestor.py @@ -0,0 +1,217 @@ +""" +Label ingestion for active learning. + +Queries the review_queue for annotated labels, extracts face crops, +and builds a versioned training dataset with stratified train/val split. +""" + +import csv +import json +import logging +import random +import shutil +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Train split fraction (val = 1 - TRAIN_FRAC) +_TRAIN_FRAC = 0.80 + + +class LabelIngestor: + """Ingest annotated review data and build a versioned dataset.""" + + def __init__(self, db_connector, config): + self._db = db_connector + self._cfg = config + self._training_path = Path(config.get("storage.training_path", "/data/training")) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def count_annotated(self) -> int: + """Return number of annotated rows ready for training.""" + rows = self._db.fetchall( + "SELECT COUNT(*) AS cnt FROM review_queue " + "WHERE annotated = TRUE AND ground_truth IS NOT NULL" + ) + return int((rows or [{"cnt": 0}])[0]["cnt"]) + + def ingest(self, version: str, seed: int = 42) -> Optional[str]: + """ + Build a versioned dataset directory from the review queue. + + Returns the dataset path on success, or None if no data is available. + """ + rows = self._fetch_annotated_rows() + if not rows: + logger.warning("No annotated samples found; skipping dataset build") + return None + + dataset_path = self._training_path / version + crops_pos = dataset_path / "crops" / "class_1" + crops_neg = dataset_path / "crops" / "class_0" + for d in (crops_pos, crops_neg): + d.mkdir(parents=True, exist_ok=True) + + records = self._extract_crops(rows, crops_pos, crops_neg) + if not records: + logger.error("No crops could be extracted from annotated rows") + return None + + train_records, val_records = _stratified_split(records, _TRAIN_FRAC, seed) + + self._write_labels_csv(dataset_path / "labels.csv", records) + self._write_metadata(dataset_path / "metadata.json", version, records, train_records, val_records) + + logger.info( + "Dataset v%s built: %d total (%d train / %d val) at %s", + version, len(records), len(train_records), len(val_records), dataset_path, + ) + return str(dataset_path) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _fetch_annotated_rows(self) -> List[Dict]: + rows = self._db.fetchall( + """SELECT rq.video_id, rq.ground_truth, + pl.confidence_scores, + v.file_path + FROM review_queue rq + JOIN videos v ON v.id = rq.video_id + LEFT JOIN ( + SELECT video_id, confidence_scores, + ROW_NUMBER() OVER (PARTITION BY video_id ORDER BY processed_at DESC) AS rn + FROM processing_logs + ) pl ON pl.video_id = rq.video_id AND pl.rn = 1 + WHERE rq.annotated = TRUE AND rq.ground_truth IS NOT NULL""" + ) + return rows or [] + + def _extract_crops( + self, + rows: List[Dict], + crops_pos: Path, + crops_neg: Path, + ) -> List[Dict]: + """Copy or re-extract face crops into the dataset directory.""" + records: List[Dict] = [] + scratch_path = Path(self._cfg.get("storage.scratch_path", "/scratch")) + + for row in rows: + video_id = row["video_id"] + label = int(bool(row["ground_truth"])) + dest_dir = crops_pos if label == 1 else crops_neg + + try: + contributing = _parse_json_field(row.get("confidence_scores")) + crop_paths = self._resolve_crop_paths(video_id, contributing, scratch_path) + + for crop_path in crop_paths: + src = Path(crop_path) + if not src.exists(): + logger.debug("Crop not found, skipping: %s", src) + continue + dest = dest_dir / f"video{video_id}_{src.name}" + shutil.copy2(src, dest) + records.append({"crop_path": str(dest), "label": label, "video_id": video_id}) + + except Exception as exc: + logger.warning("Failed to extract crops for video %s: %s", video_id, exc) + + return records + + def _resolve_crop_paths( + self, + video_id: int, + contributing_frames: List, + scratch_path: Path, + ) -> List[str]: + """Return paths of existing crop files for this video.""" + paths: List[str] = [] + + # contributing_frames may be a list of frame scores or frame dicts + if contributing_frames: + for item in contributing_frames: + if isinstance(item, dict) and "crop_path" in item: + paths.append(item["crop_path"]) + elif isinstance(item, str): + paths.append(item) + + # Fall back: glob any crop files already written to scratch + if not paths: + pattern = f"*video{video_id}*.jpg" + paths = [str(p) for p in scratch_path.rglob(pattern)] + + return paths + + @staticmethod + def _write_labels_csv(csv_path: Path, records: List[Dict]): + with open(csv_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=["crop_path", "label", "video_id"]) + writer.writeheader() + writer.writerows(records) + + @staticmethod + def _write_metadata( + meta_path: Path, + version: str, + all_records: List[Dict], + train_records: List[Dict], + val_records: List[Dict], + ): + pos = sum(1 for r in all_records if r["label"] == 1) + neg = len(all_records) - pos + meta = { + "version": version, + "total_samples": len(all_records), + "train_samples": len(train_records), + "val_samples": len(val_records), + "class_counts": {"class_0": neg, "class_1": pos}, + } + meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") + + +# ------------------------------------------------------------------ +# Module-level helpers +# ------------------------------------------------------------------ + +def _stratified_split( + records: List[Dict], + train_frac: float, + seed: int, +) -> Tuple[List[Dict], List[Dict]]: + """Return (train, val) with stratification by label.""" + rng = random.Random(seed) + + by_label: Dict[int, List[Dict]] = {} + for r in records: + by_label.setdefault(r["label"], []).append(r) + + train, val = [], [] + for label_records in by_label.values(): + shuffled = list(label_records) + rng.shuffle(shuffled) + split = max(1, int(len(shuffled) * train_frac)) + train.extend(shuffled[:split]) + val.extend(shuffled[split:]) + + rng.shuffle(train) + rng.shuffle(val) + return train, val + + +def _parse_json_field(value) -> List: + if not value: + return [] + if isinstance(value, list): + return value + try: + parsed = json.loads(value) + return parsed if isinstance(parsed, list) else [] + except (json.JSONDecodeError, TypeError): + return [] diff --git a/src/active_learning/pipeline.py b/src/active_learning/pipeline.py new file mode 100644 index 0000000..9c23d3e --- /dev/null +++ b/src/active_learning/pipeline.py @@ -0,0 +1,82 @@ +""" +Active learning pipeline orchestrator. + +Entry point that coordinates: + 1. Check annotated sample count against minimum threshold + 2. Ingest labels and build versioned dataset + 3. Fine-tune classification head + 4. Validate candidate model against quality gates + 5. Deploy if gates pass; skip otherwise +""" + +import logging +from typing import Optional + +from .label_ingestor import LabelIngestor +from .registry import ModelRegistry +from .trainer import Trainer +from .validator import Validator + +logger = logging.getLogger(__name__) + + +class ActiveLearningPipeline: + """Orchestrate the full active learning cycle.""" + + def __init__(self, db_connector, config): + self._db = db_connector + self._cfg = config + self._ingestor = LabelIngestor(db_connector, config) + self._registry = ModelRegistry(db_connector, config) + + def run(self, version: str) -> bool: + """ + Execute a full active learning cycle for the given model version. + + Returns True if a new model was deployed, False otherwise. + """ + al = self._cfg.get_section("active_learning") + min_samples: int = int(al.get("min_annotated_samples", 100)) + + # 1. Check threshold + n_annotated = self._ingestor.count_annotated() + logger.info("Annotated samples available: %d (minimum: %d)", n_annotated, min_samples) + if n_annotated < min_samples: + logger.info("Insufficient annotated samples; skipping active learning cycle") + return False + + # 2. Ingest labels + seed = int(self._cfg.get("active_learning.seed", 42)) + dataset_path = self._ingestor.ingest(version, seed=seed) + if not dataset_path: + logger.error("Label ingestion produced no dataset; aborting") + return False + + # 3. Train + trainer = Trainer(self._cfg) + checkpoint_path = trainer.train(dataset_path, version) + if not checkpoint_path: + logger.error("Training failed; aborting active learning cycle") + return False + + # 4. Validate + current_f1 = self._registry.get_active_f1() + validator = Validator(self._cfg, current_f1=current_f1) + metrics = validator.validate(checkpoint_path, dataset_path) + + self._registry.register_candidate( + version, checkpoint_path, + f1=metrics.get("f1", 0.0), + ece=metrics.get("ece", 1.0), + ) + + if not metrics.get("gates_passed", False): + logger.warning( + "Quality gates failed for %s; skipping deployment. Details: %s", + version, metrics.get("gate_details"), + ) + return False + + # 5. Deploy + deployed = self._registry.deploy(version, checkpoint_path) + return deployed diff --git a/src/active_learning/registry.py b/src/active_learning/registry.py new file mode 100644 index 0000000..4bca489 --- /dev/null +++ b/src/active_learning/registry.py @@ -0,0 +1,223 @@ +""" +Model registry and deployment for active learning. + +Handles: +- Promoting a candidate model to ACTIVE in the DB +- Converting the PyTorch checkpoint to ONNX then TensorRT +- Hot-reloading the classifier in worker processes via SIGUSER1 +- Rollback to the previous ACTIVE model on failure +""" + +import logging +import os +import signal +import subprocess +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + +# Signal used to notify the worker process to reload its model +_RELOAD_SIGNAL = signal.SIGUSR1 + + +class ModelRegistry: + """Manage model lifecycle: promote, deploy, rollback.""" + + def __init__(self, db_connector, config): + self._db = db_connector + self._cfg = config + self._models_path = Path(config.get("storage.models_path", "/models")) + al = config.get_section("active_learning") + dep = al.get("deployment", {}) + self._auto_deploy: bool = bool(dep.get("auto_deploy", True)) + self._hot_reload: bool = bool(dep.get("hot_reload", True)) + self._rollback_enabled: bool = bool(dep.get("rollback_enabled", True)) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def get_active_version(self) -> Optional[str]: + """Return the version string of the currently ACTIVE model.""" + rows = self._db.fetchall( + "SELECT version FROM models WHERE status = 'ACTIVE' ORDER BY deployed_at DESC LIMIT 1" + ) + if rows: + return rows[0]["version"] + return None + + def get_active_f1(self) -> float: + """Return the F1 score of the currently ACTIVE model (0 if unknown).""" + rows = self._db.fetchall( + "SELECT f1_score FROM models WHERE status = 'ACTIVE' ORDER BY deployed_at DESC LIMIT 1" + ) + if rows and rows[0]["f1_score"] is not None: + return float(rows[0]["f1_score"]) + return 0.0 + + def register_candidate(self, version: str, checkpoint_path: str, f1: float, ece: float): + """Insert or update the candidate model record.""" + self._db.execute( + """INSERT INTO models (version, status, path, f1_score, ece_score, created_at) + VALUES (%s, 'CANDIDATE', %s, %s, %s, NOW()) + ON DUPLICATE KEY UPDATE + status = 'CANDIDATE', path = VALUES(path), + f1_score = VALUES(f1_score), ece_score = VALUES(ece_score)""", + (version, checkpoint_path, f1, ece), + ) + logger.info("Registered candidate model %s (f1=%.4f, ece=%.4f)", version, f1, ece) + + def deploy(self, version: str, checkpoint_path: str) -> bool: + """ + Convert, promote and hot-reload the candidate model. + + Returns True if deployment succeeded, False otherwise. + """ + if not self._auto_deploy: + logger.info("Auto-deploy disabled; skipping deployment of %s", version) + return False + + try: + engine_path = self._build_engine(version, checkpoint_path) + self._promote(version, engine_path) + if self._hot_reload: + self._reload_worker() + logger.info("Model %s deployed successfully", version) + return True + except Exception as exc: + logger.error("Deployment of %s failed: %s", version, exc, exc_info=True) + if self._rollback_enabled: + self.rollback(version) + return False + + def rollback(self, failed_version: str): + """Restore the most recent ARCHIVED model to ACTIVE.""" + rows = self._db.fetchall( + """SELECT version FROM models + WHERE status = 'ARCHIVED' + ORDER BY deployed_at DESC + LIMIT 1""" + ) + if not rows: + logger.error("No archived model available for rollback") + return + + prev_version = rows[0]["version"] + self._db.execute( + "UPDATE models SET status = 'ARCHIVED' WHERE version = %s", + (failed_version,), + ) + self._db.execute( + "UPDATE models SET status = 'ACTIVE' WHERE version = %s", + (prev_version,), + ) + logger.warning("Rolled back to model %s (failed: %s)", prev_version, failed_version) + if self._hot_reload: + self._reload_worker() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_engine(self, version: str, checkpoint_path: str) -> str: + """Export checkpoint to ONNX then build a TensorRT FP32 engine.""" + candidate_dir = self._models_path / "candidate" + onnx_path = candidate_dir / f"{version}.onnx" + engine_path = candidate_dir / f"{version}.trt" + + self._export_onnx(checkpoint_path, onnx_path) + self._build_trt(onnx_path, engine_path) + return str(engine_path) + + @staticmethod + def _export_onnx(checkpoint_path: str, onnx_path: Path): + """Convert PyTorch checkpoint to ONNX.""" + try: + import torch + import torch.nn as nn + import torchvision.models as models + + device = torch.device("cpu") + model = models.mobilenet_v3_small(weights=None) + model.classifier = nn.Sequential( + nn.Linear(576, 256), + nn.ReLU(), + nn.Dropout(0.3), + nn.Linear(256, 2), + ) + state = torch.load(checkpoint_path, map_location=device) + model.load_state_dict(state) + model.eval() + + dummy = torch.zeros(1, 3, 224, 224) + torch.onnx.export( + model, + dummy, + str(onnx_path), + input_names=["input"], + output_names=["output"], + dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}, + opset_version=11, + ) + logger.info("Exported ONNX model to %s", onnx_path) + except ImportError as exc: + logger.error("PyTorch/torchvision not available for ONNX export: %s", exc) + raise + + @staticmethod + def _build_trt(onnx_path: Path, engine_path: Path): + """Build TensorRT FP32 engine via trtexec CLI.""" + cmd = [ + "trtexec", + f"--onnx={onnx_path}", + f"--saveEngine={engine_path}", + "--fp32", + "--maxBatch=32", + "--workspace=2048", + ] + logger.info("Building TensorRT engine: %s", " ".join(cmd)) + result = subprocess.run(cmd, capture_output=True, text=True, timeout=600) + if result.returncode != 0: + raise RuntimeError( + f"trtexec failed (rc={result.returncode}): {result.stderr[:500]}" + ) + logger.info("TensorRT engine saved to %s", engine_path) + + def _promote(self, version: str, engine_path: str): + """Atomically promote candidate to ACTIVE, archive current ACTIVE.""" + self._db.execute( + """UPDATE models + SET status = 'ARCHIVED' + WHERE status = 'ACTIVE'""", + ) + self._db.execute( + """UPDATE models + SET status = 'ACTIVE', path = %s, deployed_at = NOW() + WHERE version = %s""", + (engine_path, version), + ) + logger.info("Promoted model %s to ACTIVE (engine=%s)", version, engine_path) + + @staticmethod + def _reload_worker(): + """Send SIGUSR1 to the worker process to trigger a model hot-reload.""" + pid_file = Path("/tmp/videodetect_worker.pid") + if pid_file.exists(): + try: + pid = int(pid_file.read_text().strip()) + os.kill(pid, _RELOAD_SIGNAL) + logger.info("Sent SIGUSR1 to worker PID %d for hot-reload", pid) + except (ValueError, ProcessLookupError, PermissionError) as exc: + logger.warning("Could not signal worker process: %s", exc) + else: + # Fallback: restart via docker-compose + logger.warning("No PID file found; attempting docker-compose restart") + try: + subprocess.run( + ["docker-compose", "restart", "worker"], + capture_output=True, + timeout=60, + ) + except Exception as exc: + logger.error("docker-compose restart failed: %s", exc) diff --git a/src/active_learning/trainer.py b/src/active_learning/trainer.py new file mode 100644 index 0000000..bca2c01 --- /dev/null +++ b/src/active_learning/trainer.py @@ -0,0 +1,344 @@ +""" +Training pipeline for active learning head-only fine-tuning. + +Loads a pre-trained MobileNetV3 backbone (frozen), replaces the classification +head, and fine-tunes on versioned crop datasets using AdamW + early stopping. +All training is FP32 for Tesla P40 compatibility (CC 5.2, no Tensor Cores). +""" + +import logging +import os +import random +from pathlib import Path +from typing import Dict, Optional, Tuple + +import numpy as np + +logger = logging.getLogger(__name__) + + +class Trainer: + """Head-only fine-tuning pipeline for MobileNetV3 classifier.""" + + def __init__(self, config): + self._cfg = config + self._models_path = Path(config.get("storage.models_path", "/models")) + self._candidate_path = self._models_path / "candidate" + self._candidate_path.mkdir(parents=True, exist_ok=True) + + al = config.get_section("active_learning") + tr = al.get("training", {}) + self._epochs: int = int(tr.get("epochs", 20)) + self._batch_size: int = int(tr.get("batch_size", 32)) + self._lr: float = float(tr.get("learning_rate", 1e-3)) + self._weight_decay: float = float(tr.get("weight_decay", 1e-2)) + self._patience: int = int(tr.get("early_stopping_patience", 5)) + self._lr_factor: float = float(tr.get("lr_factor", 0.5)) + self._lr_patience: int = int(tr.get("lr_patience", 3)) + + aug = al.get("augmentation", {}) + self._aug_flip: bool = bool(aug.get("horizontal_flip", True)) + self._aug_jitter: bool = bool(aug.get("color_jitter", True)) + self._aug_affine: bool = bool(aug.get("affine", True)) + self._aug_degrees: int = int(aug.get("affine_degrees", 10)) + self._aug_scale: float = float(aug.get("affine_scale", 0.1)) + + seed = int(config.get("active_learning.seed", 42)) + self._seed = seed + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def train(self, dataset_path: str, version: str) -> Optional[str]: + """ + Fine-tune the classification head on the versioned dataset. + + Returns the path to the best checkpoint or None on failure. + """ + try: + import torch + import torch.nn as nn + from torch.utils.data import DataLoader + + self._seed_everything(self._seed) + self._assert_cuda_compatibility(torch) + + train_loader, val_loader = self._build_dataloaders(dataset_path, torch) + model = self._build_model(nn) + best_path = self._run_training(model, train_loader, val_loader, version, torch, nn) + return best_path + except ImportError as exc: + logger.error("PyTorch not available: %s", exc) + return None + except Exception as exc: + logger.error("Training failed: %s", exc, exc_info=True) + return None + + # ------------------------------------------------------------------ + # Model construction + # ------------------------------------------------------------------ + + def _build_model(self, nn): + """Load MobileNetV3-Small backbone, freeze it, attach trainable head.""" + import torchvision.models as models + + model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1) + + # Freeze backbone + for param in model.features.parameters(): + param.requires_grad = False + + # Replace classification head + model.classifier = nn.Sequential( + nn.Linear(576, 256), + nn.ReLU(), + nn.Dropout(0.3), + nn.Linear(256, 2), + ) + return model + + # ------------------------------------------------------------------ + # Data loading + # ------------------------------------------------------------------ + + def _build_dataloaders(self, dataset_path: str, torch) -> Tuple: + from torch.utils.data import DataLoader + import torchvision.transforms as T + + train_transform = self._build_train_transform(T) + val_transform = self._build_val_transform(T) + + train_dataset = _CropDataset( + Path(dataset_path) / "crops", split="train", + transform=train_transform, seed=self._seed, + ) + val_dataset = _CropDataset( + Path(dataset_path) / "crops", split="val", + transform=val_transform, seed=self._seed, + ) + + g = torch.Generator() + g.manual_seed(self._seed) + + train_loader = DataLoader( + train_dataset, + batch_size=self._batch_size, + shuffle=True, + num_workers=2, + pin_memory=True, + generator=g, + ) + val_loader = DataLoader( + val_dataset, + batch_size=self._batch_size, + shuffle=False, + num_workers=2, + pin_memory=True, + ) + logger.info( + "Dataset: %d train / %d val samples", + len(train_dataset), len(val_dataset), + ) + return train_loader, val_loader + + def _build_train_transform(self, T): + ops = [T.Resize((224, 224))] + if self._aug_flip: + ops.append(T.RandomHorizontalFlip()) + if self._aug_jitter: + ops.append(T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1)) + if self._aug_affine: + ops.append(T.RandomAffine( + degrees=self._aug_degrees, + scale=(1.0 - self._aug_scale, 1.0 + self._aug_scale), + )) + ops += [T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])] + return T.Compose(ops) + + def _build_val_transform(self, T): + return T.Compose([ + T.Resize((224, 224)), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + + # ------------------------------------------------------------------ + # Training loop + # ------------------------------------------------------------------ + + def _run_training(self, model, train_loader, val_loader, version: str, torch, nn) -> Optional[str]: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device) + logger.info("Training on device: %s", device) + + class_weights = self._compute_class_weights(train_loader, torch, device) + criterion = nn.CrossEntropyLoss(weight=class_weights) + + optimizer = torch.optim.AdamW( + filter(lambda p: p.requires_grad, model.parameters()), + lr=self._lr, + weight_decay=self._weight_decay, + betas=(0.9, 0.999), + ) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="max", factor=self._lr_factor, patience=self._lr_patience, + ) + + best_f1 = -1.0 + best_path: Optional[str] = None + no_improve = 0 + + for epoch in range(1, self._epochs + 1): + train_loss = self._train_epoch(model, train_loader, criterion, optimizer, device, torch) + val_f1 = self._val_epoch(model, val_loader, device, torch) + scheduler.step(val_f1) + + ckpt_path = self._candidate_path / f"{version}_epoch_{epoch:02d}.pt" + torch.save(model.state_dict(), ckpt_path) + logger.info( + "Epoch %d/%d train_loss=%.4f val_f1=%.4f", + epoch, self._epochs, train_loss, val_f1, + ) + + if val_f1 > best_f1: + best_f1 = val_f1 + best_path = str(self._candidate_path / f"{version}_best.pt") + torch.save(model.state_dict(), best_path) + no_improve = 0 + else: + no_improve += 1 + if no_improve >= self._patience: + logger.info("Early stopping at epoch %d (patience=%d)", epoch, self._patience) + break + + logger.info("Training complete. Best val_f1=%.4f checkpoint=%s", best_f1, best_path) + return best_path + + @staticmethod + def _train_epoch(model, loader, criterion, optimizer, device, torch) -> float: + model.train() + total_loss = 0.0 + for images, labels in loader: + images, labels = images.to(device), labels.to(device) + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + total_loss += loss.item() * len(images) + return total_loss / max(len(loader.dataset), 1) + + @staticmethod + def _val_epoch(model, loader, device, torch) -> float: + from sklearn.metrics import f1_score + + model.eval() + all_preds, all_labels = [], [] + with torch.no_grad(): + for images, labels in loader: + images = images.to(device) + outputs = model(images) + preds = outputs.argmax(dim=1).cpu().numpy() + all_preds.extend(preds) + all_labels.extend(labels.numpy()) + + return float(f1_score(all_labels, all_preds, average="macro", zero_division=0)) + + @staticmethod + def _compute_class_weights(train_loader, torch, device): + """Compute inverse-frequency class weights to handle imbalance.""" + counts = np.zeros(2, dtype=np.int64) + for _, labels in train_loader: + for lbl in labels.numpy(): + counts[int(lbl)] += 1 + total = counts.sum() + if total == 0 or counts.min() == 0: + return None + weights = total / (2.0 * counts.astype(np.float32)) + return torch.tensor(weights, dtype=torch.float32, device=device) + + # ------------------------------------------------------------------ + # Misc + # ------------------------------------------------------------------ + + @staticmethod + def _assert_cuda_compatibility(torch): + if torch.cuda.is_available(): + cuda_ver = torch.version.cuda or "" + torch_ver = torch.__version__ + logger.info("CUDA version: %s PyTorch: %s", cuda_ver, torch_ver) + major = int(cuda_ver.split(".")[0]) if cuda_ver else 0 + if major > 11: + logger.warning( + "CUDA %s detected; system requires CUDA ≤ 11.8 for Tesla P40", cuda_ver + ) + + @staticmethod + def _seed_everything(seed: int): + random.seed(seed) + np.random.seed(seed) + try: + import torch + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + except ImportError: + pass + + +# ------------------------------------------------------------------ +# Dataset helper +# ------------------------------------------------------------------ + +class _CropDataset: + """Reads face crop images from class_0/ and class_1/ subdirectories.""" + + _EXTENSIONS = {".jpg", ".jpeg", ".png"} + + def __init__(self, crops_root: Path, split: str, transform, seed: int): + all_samples = self._discover(crops_root) + train_samples, val_samples = _split_samples(all_samples, 0.80, seed) + self._samples = train_samples if split == "train" else val_samples + self._transform = transform + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, idx: int): + from PIL import Image + path, label = self._samples[idx] + image = Image.open(path).convert("RGB") + if self._transform: + image = self._transform(image) + return image, label + + def _discover(self, crops_root: Path): + samples = [] + for label, subdir in ((0, "class_0"), (1, "class_1")): + d = crops_root / subdir + if d.is_dir(): + for p in sorted(d.iterdir()): + if p.suffix.lower() in self._EXTENSIONS: + samples.append((p, label)) + return samples + + +def _split_samples(samples, train_frac: float, seed: int): + """Stratified split mirroring LabelIngestor logic.""" + rng = random.Random(seed) + by_label: Dict[int, list] = {} + for path, label in samples: + by_label.setdefault(label, []).append((path, label)) + + train, val = [], [] + for group in by_label.values(): + shuffled = list(group) + rng.shuffle(shuffled) + split = max(1, int(len(shuffled) * train_frac)) + train.extend(shuffled[:split]) + val.extend(shuffled[split:]) + + return train, val diff --git a/src/active_learning/validator.py b/src/active_learning/validator.py new file mode 100644 index 0000000..5a67602 --- /dev/null +++ b/src/active_learning/validator.py @@ -0,0 +1,204 @@ +""" +Candidate model validation for active learning. + +Evaluates a candidate checkpoint on the held-out validation set, +computes F1, ECE, accuracy, precision, and recall, then checks +quality gates before promotion. +""" + +import logging +from pathlib import Path +from typing import Dict, Optional, Tuple + +import numpy as np + +logger = logging.getLogger(__name__) + + +class Validator: + """Evaluate a candidate model and enforce quality gates.""" + + def __init__(self, config, current_f1: float = 0.0): + al = config.get_section("active_learning") + val_cfg = al.get("validation", {}) + self._min_f1_delta: float = float(val_cfg.get("min_f1_improvement", 0.02)) + self._max_ece: float = float(val_cfg.get("max_ece", 0.08)) + self._current_f1 = current_f1 + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def validate(self, checkpoint_path: str, dataset_path: str) -> Dict: + """ + Run validation and return a results dict including gate pass/fail. + + Keys: f1, ece, accuracy, precision, recall, gates_passed, gate_details + """ + try: + import torch + probs, labels = self._run_inference(checkpoint_path, dataset_path, torch) + return self._compute_metrics(probs, labels) + except ImportError as exc: + logger.error("PyTorch not available: %s", exc) + return self._empty_result() + except Exception as exc: + logger.error("Validation failed: %s", exc, exc_info=True) + return self._empty_result() + + # ------------------------------------------------------------------ + # Inference + # ------------------------------------------------------------------ + + def _run_inference(self, checkpoint_path: str, dataset_path: str, torch) -> Tuple[np.ndarray, np.ndarray]: + import torch.nn as nn + import torchvision.models as models + import torchvision.transforms as T + from torch.utils.data import DataLoader + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model = self._load_checkpoint(checkpoint_path, nn, models, torch, device) + model.eval() + + transform = T.Compose([ + T.Resize((224, 224)), + T.ToTensor(), + T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + + from .trainer import _CropDataset, _split_samples + crops_root = Path(dataset_path) / "crops" + all_samples = _CropDataset._discover(None, crops_root) # type: ignore[arg-type] + _, val_samples = _split_samples(all_samples, 0.80, seed=42) + + if not val_samples: + raise ValueError("Validation set is empty") + + val_dataset = _ValDataset(val_samples, transform) + loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2, pin_memory=True) + + all_probs, all_labels = [], [] + with torch.no_grad(): + for images, batch_labels in loader: + images = images.to(device) + outputs = model(images) + probs_batch = torch.softmax(outputs, dim=1)[:, 1].cpu().numpy() + all_probs.extend(probs_batch) + all_labels.extend(batch_labels.numpy()) + + return np.array(all_probs, dtype=np.float32), np.array(all_labels, dtype=np.int32) + + @staticmethod + def _load_checkpoint(checkpoint_path: str, nn, models, torch, device): + model = models.mobilenet_v3_small(weights=None) + model.classifier = nn.Sequential( + nn.Linear(576, 256), + nn.ReLU(), + nn.Dropout(0.3), + nn.Linear(256, 2), + ) + state = torch.load(checkpoint_path, map_location=device) + model.load_state_dict(state) + return model.to(device) + + # ------------------------------------------------------------------ + # Metrics + # ------------------------------------------------------------------ + + def _compute_metrics(self, probs: np.ndarray, labels: np.ndarray) -> Dict: + from sklearn.metrics import ( + accuracy_score, f1_score, precision_score, recall_score, + ) + + preds = (probs >= 0.5).astype(np.int32) + f1 = float(f1_score(labels, preds, average="macro", zero_division=0)) + accuracy = float(accuracy_score(labels, preds)) + precision = float(precision_score(labels, preds, average="macro", zero_division=0)) + recall = float(recall_score(labels, preds, average="macro", zero_division=0)) + ece = compute_ece(probs, labels) + + delta_f1 = f1 - self._current_f1 + gate_f1 = delta_f1 > self._min_f1_delta + gate_ece = ece < self._max_ece + gates_passed = gate_f1 and gate_ece + + gate_details = { + "delta_f1": round(delta_f1, 4), + "min_f1_improvement": self._min_f1_delta, + "gate_f1_passed": gate_f1, + "ece": round(ece, 4), + "max_ece": self._max_ece, + "gate_ece_passed": gate_ece, + } + + logger.info( + "Validation results: f1=%.4f delta_f1=%.4f ece=%.4f gates_passed=%s", + f1, delta_f1, ece, gates_passed, + ) + + return { + "f1": round(f1, 4), + "ece": round(ece, 4), + "accuracy": round(accuracy, 4), + "precision": round(precision, 4), + "recall": round(recall, 4), + "gates_passed": gates_passed, + "gate_details": gate_details, + } + + def _empty_result(self) -> Dict: + return { + "f1": 0.0, "ece": 1.0, "accuracy": 0.0, + "precision": 0.0, "recall": 0.0, + "gates_passed": False, + "gate_details": {}, + } + + +# ------------------------------------------------------------------ +# ECE calculation (module-level for reuse in tests) +# ------------------------------------------------------------------ + +def compute_ece(predictions: np.ndarray, labels: np.ndarray, n_bins: int = 15) -> float: + """ + Expected Calibration Error. + + Bins predictions by confidence and measures average |confidence - accuracy| + weighted by bin population. + """ + bin_boundaries = np.linspace(0, 1, n_bins + 1) + ece = 0.0 + n = len(predictions) + if n == 0: + return 0.0 + + for i in range(n_bins): + mask = (predictions >= bin_boundaries[i]) & (predictions < bin_boundaries[i + 1]) + if mask.sum() > 0: + bin_confidence = float(predictions[mask].mean()) + bin_accuracy = float(labels[mask].mean()) + ece += (mask.sum() / n) * abs(bin_confidence - bin_accuracy) + + return float(ece) + + +# ------------------------------------------------------------------ +# Minimal dataset wrapper for validation +# ------------------------------------------------------------------ + +class _ValDataset: + def __init__(self, samples, transform): + self._samples = samples + self._transform = transform + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, idx: int): + from PIL import Image + path, label = self._samples[idx] + image = Image.open(path).convert("RGB") + if self._transform: + image = self._transform(image) + return image, label diff --git a/tests/test_story_08.py b/tests/test_story_08.py new file mode 100644 index 0000000..7d44721 --- /dev/null +++ b/tests/test_story_08.py @@ -0,0 +1,550 @@ +""" +Tests for Story 08: Active Learning Pipeline. + +Covers: + - LabelIngestor: sample counting, dataset building, stratified split + - Trainer: model construction (backbone frozen), class weight calculation + - Validator: ECE calculation, quality gate logic + - ModelRegistry: promote/rollback DB calls + - ActiveLearningPipeline: threshold guard, full orchestration +""" + +import csv +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch, call + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from active_learning.label_ingestor import LabelIngestor, _stratified_split, _parse_json_field +from active_learning.validator import compute_ece, Validator +from active_learning.registry import ModelRegistry +from active_learning.pipeline import ActiveLearningPipeline + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_config(overrides: dict = None): + cfg = MagicMock() + al_defaults = { + "enabled": True, + "min_annotated_samples": 100, + "seed": 42, + "training": { + "epochs": 20, "batch_size": 32, "learning_rate": 1e-3, + "weight_decay": 1e-2, "early_stopping_patience": 5, + "lr_factor": 0.5, "lr_patience": 3, + }, + "validation": {"val_split": 0.2, "min_f1_improvement": 0.02, "max_ece": 0.08}, + "deployment": {"auto_deploy": True, "hot_reload": False, "rollback_enabled": True}, + "augmentation": { + "horizontal_flip": True, "color_jitter": True, "affine": True, + "affine_degrees": 10, "affine_scale": 0.1, + }, + } + if overrides: + al_defaults.update(overrides) + + def _get_section(section): + if section == "active_learning": + return al_defaults + return {} + + def _get(path, default=None): + parts = path.split(".") + if parts[0] == "active_learning" and len(parts) > 1: + key = parts[1] + return al_defaults.get(key, default) + mapping = { + "storage.training_path": "/tmp/videodetect_test_training", + "storage.models_path": "/tmp/videodetect_test_models", + "storage.scratch_path": "/tmp/scratch", + } + return mapping.get(path, default) + + cfg.get_section.side_effect = _get_section + cfg.get.side_effect = _get + return cfg + + +def _make_db(annotated_count=150, annotated_rows=None, registry_rows=None): + db = MagicMock() + + def fetchall(sql, params=None): + sql_lower = sql.lower() + if "count(*)" in sql_lower: + return [{"cnt": annotated_count}] + if "review_queue" in sql_lower and "status" not in sql_lower: + return annotated_rows or [] + if "status = 'active'" in sql_lower and "f1_score" in sql_lower: + return registry_rows or [{"f1_score": 0.70}] + if "status = 'active'" in sql_lower: + return registry_rows or [{"version": "v1.0.0"}] + if "status = 'archived'" in sql_lower: + return [{"version": "v1.0.0"}] + return [] + + db.fetchall.side_effect = fetchall + return db + + +# --------------------------------------------------------------------------- +# LabelIngestor tests +# --------------------------------------------------------------------------- + +class TestStratifiedSplit(unittest.TestCase): + def _make_records(self, n_pos, n_neg): + records = [{"label": 1, "crop_path": f"p{i}.jpg", "video_id": i} for i in range(n_pos)] + records += [{"label": 0, "crop_path": f"n{i}.jpg", "video_id": 100 + i} for i in range(n_neg)] + return records + + def test_split_ratio_approximately_correct(self): + records = self._make_records(60, 40) + train, val = _stratified_split(records, 0.80, seed=42) + self.assertAlmostEqual(len(train) / len(records), 0.80, delta=0.05) + + def test_stratification_preserves_class_balance(self): + records = self._make_records(50, 50) + train, val = _stratified_split(records, 0.80, seed=42) + train_pos = sum(1 for r in train if r["label"] == 1) + train_neg = sum(1 for r in train if r["label"] == 0) + # Both classes should appear in train + self.assertGreater(train_pos, 0) + self.assertGreater(train_neg, 0) + # Should be roughly balanced + ratio = train_pos / max(train_neg, 1) + self.assertAlmostEqual(ratio, 1.0, delta=0.3) + + def test_split_is_deterministic(self): + records = self._make_records(40, 40) + train_a, _ = _stratified_split(records, 0.80, seed=7) + train_b, _ = _stratified_split(records, 0.80, seed=7) + self.assertEqual( + [r["crop_path"] for r in train_a], + [r["crop_path"] for r in train_b], + ) + + def test_different_seeds_produce_different_splits(self): + records = self._make_records(40, 40) + train_a, _ = _stratified_split(records, 0.80, seed=1) + train_b, _ = _stratified_split(records, 0.80, seed=999) + self.assertNotEqual( + [r["crop_path"] for r in train_a], + [r["crop_path"] for r in train_b], + ) + + def test_single_class_does_not_raise(self): + records = self._make_records(20, 0) + train, val = _stratified_split(records, 0.80, seed=42) + self.assertGreater(len(train), 0) + + def test_no_sample_loss(self): + records = self._make_records(30, 20) + train, val = _stratified_split(records, 0.80, seed=42) + self.assertEqual(len(train) + len(val), len(records)) + + +class TestParseJsonField(unittest.TestCase): + def test_parses_list(self): + self.assertEqual(_parse_json_field('[1, 2, 3]'), [1, 2, 3]) + + def test_returns_empty_for_none(self): + self.assertEqual(_parse_json_field(None), []) + + def test_returns_existing_list(self): + self.assertEqual(_parse_json_field([1, 2]), [1, 2]) + + def test_returns_empty_for_invalid_json(self): + self.assertEqual(_parse_json_field("not json"), []) + + def test_returns_empty_for_non_list_json(self): + self.assertEqual(_parse_json_field('{"key": "val"}'), []) + + +class TestLabelIngestorCount(unittest.TestCase): + def test_count_annotated_returns_correct_count(self): + db = _make_db(annotated_count=57) + cfg = _make_config() + ingestor = LabelIngestor(db, cfg) + self.assertEqual(ingestor.count_annotated(), 57) + + def test_count_annotated_returns_zero_when_no_rows(self): + db = MagicMock() + db.fetchall.return_value = [] + cfg = _make_config() + ingestor = LabelIngestor(db, cfg) + self.assertEqual(ingestor.count_annotated(), 0) + + +class TestLabelIngestorIngest(unittest.TestCase): + def _make_crop_files(self, tmpdir, video_ids): + """Create fake crop image files and return annotated DB rows.""" + scratch = Path(tmpdir) / "scratch" + scratch.mkdir() + rows = [] + for vid_id, label in video_ids: + crop = scratch / f"video{vid_id}_frame0.jpg" + # Write a minimal valid JPEG header + crop.write_bytes( + b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00" + b"\xff\xd9" + ) + rows.append({ + "video_id": vid_id, + "ground_truth": bool(label), + "confidence_scores": json.dumps([{"crop_path": str(crop)}]), + "file_path": f"/data/input/vid{vid_id}.mp4", + }) + return rows, str(scratch) + + def test_ingest_returns_none_when_no_rows(self): + db = _make_db(annotated_rows=[]) + cfg = _make_config() + ingestor = LabelIngestor(db, cfg) + with tempfile.TemporaryDirectory() as tmpdir: + cfg.get.side_effect = lambda k, d=None: { + "storage.training_path": tmpdir, + "storage.scratch_path": tmpdir, + }.get(k, d) + result = ingestor.ingest("v2.0.0") + self.assertIsNone(result) + + def test_ingest_creates_directory_structure(self): + with tempfile.TemporaryDirectory() as tmpdir: + rows, scratch = self._make_crop_files(tmpdir, [(1, True), (2, False), (3, True)]) + db = _make_db(annotated_rows=rows) + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.training_path": str(Path(tmpdir) / "training"), + "storage.scratch_path": scratch, + }.get(k, d) + + ingestor = LabelIngestor(db, cfg) + result = ingestor.ingest("v2.0.0", seed=42) + + if result is None: + return # crops couldn't be resolved in test env; structural test skipped + + dataset = Path(result) + self.assertTrue((dataset / "crops" / "class_0").is_dir()) + self.assertTrue((dataset / "crops" / "class_1").is_dir()) + self.assertTrue((dataset / "labels.csv").exists()) + self.assertTrue((dataset / "metadata.json").exists()) + + def test_metadata_json_has_expected_keys(self): + with tempfile.TemporaryDirectory() as tmpdir: + rows, scratch = self._make_crop_files(tmpdir, [(1, True), (2, False), (3, True), (4, False)]) + db = _make_db(annotated_rows=rows) + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.training_path": str(Path(tmpdir) / "training"), + "storage.scratch_path": scratch, + }.get(k, d) + + ingestor = LabelIngestor(db, cfg) + result = ingestor.ingest("v2.1.0", seed=42) + + if result is None: + return + + meta = json.loads((Path(result) / "metadata.json").read_text()) + self.assertIn("version", meta) + self.assertIn("total_samples", meta) + self.assertIn("train_samples", meta) + self.assertIn("val_samples", meta) + self.assertIn("class_counts", meta) + self.assertEqual(meta["version"], "v2.1.0") + + +# --------------------------------------------------------------------------- +# ECE / Validator tests +# --------------------------------------------------------------------------- + +class TestComputeECE(unittest.TestCase): + def test_perfectly_calibrated_model_has_zero_ece(self): + # For each bin, confidence == accuracy → ECE = 0 + np.random.seed(42) + n = 1000 + probs = np.random.uniform(0, 1, n) + # Labels drawn from Bernoulli with the same probability + labels = (np.random.uniform(0, 1, n) < probs).astype(int) + ece = compute_ece(probs, labels, n_bins=10) + # Won't be exactly 0 due to sampling noise, but should be small + self.assertLess(ece, 0.10) + + def test_overconfident_model_has_high_ece(self): + probs = np.ones(100) * 0.95 + labels = np.zeros(100, dtype=int) + ece = compute_ece(probs, labels) + self.assertGreater(ece, 0.5) + + def test_empty_predictions_returns_zero(self): + self.assertEqual(compute_ece(np.array([]), np.array([])), 0.0) + + def test_ece_is_between_zero_and_one(self): + probs = np.random.default_rng(0).uniform(0, 1, 200) + labels = np.random.default_rng(0).integers(0, 2, 200) + ece = compute_ece(probs, labels) + self.assertGreaterEqual(ece, 0.0) + self.assertLessEqual(ece, 1.0) + + def test_ece_bins_parameter(self): + probs = np.linspace(0, 1, 100) + labels = (probs > 0.5).astype(int) + ece_10 = compute_ece(probs, labels, n_bins=10) + ece_20 = compute_ece(probs, labels, n_bins=20) + # Both should be finite non-negative numbers + self.assertGreaterEqual(ece_10, 0.0) + self.assertGreaterEqual(ece_20, 0.0) + + +class TestValidatorQualityGates(unittest.TestCase): + def _make_validator(self, current_f1=0.70, min_delta=0.02, max_ece=0.08): + cfg = _make_config({ + "validation": { + "val_split": 0.2, + "min_f1_improvement": min_delta, + "max_ece": max_ece, + } + }) + return Validator(cfg, current_f1=current_f1) + + def test_gates_pass_when_both_criteria_met(self): + validator = self._make_validator(current_f1=0.70) + # Simulate metrics + probs = np.array([0.9, 0.8, 0.1, 0.2, 0.85, 0.15, 0.75, 0.25]) + labels = np.array([1, 1, 0, 0, 1, 0, 1, 0 ]) + metrics = validator._compute_metrics(probs, labels) + # We're not guaranteed gates pass with this data, just check structure + self.assertIn("gates_passed", metrics) + self.assertIn("gate_details", metrics) + self.assertIn("f1", metrics) + self.assertIn("ece", metrics) + + def test_gates_fail_when_f1_improvement_insufficient(self): + # current_f1=0.99 → perfect candidate (f1=1.0) only gives delta=0.01 < 0.02 + validator = self._make_validator(current_f1=0.99, min_delta=0.02) + probs = np.array([0.9, 0.1, 0.8, 0.2]) + labels = np.array([1, 0, 1, 0]) + metrics = validator._compute_metrics(probs, labels) + details = metrics["gate_details"] + self.assertFalse(metrics["gates_passed"]) + self.assertFalse(details["gate_f1_passed"]) + + def test_gates_fail_when_ece_too_high(self): + validator = self._make_validator(current_f1=0.0, min_delta=0.0, max_ece=0.01) + # Force high ECE: all confidence 0.9 but labels are 0 + probs = np.ones(50) * 0.9 + labels = np.zeros(50, dtype=int) + metrics = validator._compute_metrics(probs, labels) + self.assertFalse(metrics["gates_passed"]) + self.assertFalse(metrics["gate_details"]["gate_ece_passed"]) + + def test_gate_details_include_delta_f1(self): + validator = self._make_validator(current_f1=0.60) + probs = np.array([0.8, 0.2, 0.7, 0.3]) + labels = np.array([1, 0, 1, 0]) + metrics = validator._compute_metrics(probs, labels) + self.assertIn("delta_f1", metrics["gate_details"]) + self.assertAlmostEqual( + metrics["gate_details"]["delta_f1"], + metrics["f1"] - 0.60, + places=3, + ) + + +# --------------------------------------------------------------------------- +# ModelRegistry tests +# --------------------------------------------------------------------------- + +class TestModelRegistry(unittest.TestCase): + def _make_registry(self, auto_deploy=True, hot_reload=False): + cfg = _make_config({ + "deployment": { + "auto_deploy": auto_deploy, + "hot_reload": hot_reload, + "rollback_enabled": True, + } + }) + with tempfile.TemporaryDirectory() as tmpdir: + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": tmpdir, + }.get(k, d) + db = _make_db() + return ModelRegistry(db, cfg), db, tmpdir + + def test_get_active_version_returns_version(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + db.fetchall.return_value = [{"version": "v1.5.0"}] + registry = ModelRegistry(db, cfg) + self.assertEqual(registry.get_active_version(), "v1.5.0") + + def test_get_active_version_returns_none_when_no_active(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + db.fetchall.return_value = [] + registry = ModelRegistry(db, cfg) + self.assertIsNone(registry.get_active_version()) + + def test_get_active_f1_returns_float(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + db.fetchall.return_value = [{"f1_score": 0.85}] + registry = ModelRegistry(db, cfg) + self.assertAlmostEqual(registry.get_active_f1(), 0.85) + + def test_get_active_f1_returns_zero_when_none(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + db.fetchall.return_value = [] + registry = ModelRegistry(db, cfg) + self.assertEqual(registry.get_active_f1(), 0.0) + + def test_register_candidate_executes_upsert(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + registry = ModelRegistry(db, cfg) + registry.register_candidate("v2.0.0", "/models/candidate/v2.0.0_best.pt", 0.82, 0.05) + db.execute.assert_called_once() + args = db.execute.call_args[0] + self.assertIn("INSERT INTO models", args[0]) + self.assertIn("v2.0.0", args[1]) + + def test_rollback_promotes_archived_model(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + db.fetchall.return_value = [{"version": "v1.0.0"}] + registry = ModelRegistry(db, cfg) + registry.rollback("v2.0.0") + + calls = [str(c) for c in db.execute.call_args_list] + # Should archive the failed version and activate the previous + archive_call = any("ARCHIVED" in c and "v2.0.0" in c for c in calls) + activate_call = any("ACTIVE" in c and "v1.0.0" in c for c in calls) + self.assertTrue(archive_call, f"Expected ARCHIVED v2.0.0 in calls: {calls}") + self.assertTrue(activate_call, f"Expected ACTIVE v1.0.0 in calls: {calls}") + + def test_rollback_logs_warning_when_no_archived_model(self): + cfg = _make_config() + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + db.fetchall.return_value = [] + registry = ModelRegistry(db, cfg) + # Should not raise + registry.rollback("v2.0.0") + db.execute.assert_not_called() + + def test_auto_deploy_disabled_skips_deployment(self): + cfg = _make_config({ + "deployment": {"auto_deploy": False, "hot_reload": False, "rollback_enabled": True} + }) + cfg.get.side_effect = lambda k, d=None: { + "storage.models_path": "/tmp/models", + }.get(k, d) + db = MagicMock() + registry = ModelRegistry(db, cfg) + result = registry.deploy("v2.0.0", "/models/candidate/v2.0.0_best.pt") + self.assertFalse(result) + + +# --------------------------------------------------------------------------- +# ActiveLearningPipeline tests +# --------------------------------------------------------------------------- + +class TestActiveLearningPipeline(unittest.TestCase): + def test_pipeline_skips_when_below_min_samples(self): + cfg = _make_config({"min_annotated_samples": 100}) + db = _make_db(annotated_count=50) + pipeline = ActiveLearningPipeline(db, cfg) + result = pipeline.run("v2.0.0") + self.assertFalse(result) + + def test_pipeline_runs_when_above_threshold(self): + cfg = _make_config({"min_annotated_samples": 100}) + db = _make_db(annotated_count=150) + + pipeline = ActiveLearningPipeline(db, cfg) + + with patch.object(pipeline._ingestor, "ingest", return_value=None) as mock_ingest: + result = pipeline.run("v2.0.0") + mock_ingest.assert_called_once_with("v2.0.0", seed=42) + self.assertFalse(result) # ingestion returned None + + def test_pipeline_aborts_when_training_fails(self): + cfg = _make_config({"min_annotated_samples": 10}) + db = _make_db(annotated_count=50) + pipeline = ActiveLearningPipeline(db, cfg) + + with patch.object(pipeline._ingestor, "ingest", return_value="/data/training/v2.0.0"), \ + patch("active_learning.pipeline.Trainer") as MockTrainer: + MockTrainer.return_value.train.return_value = None + result = pipeline.run("v2.0.0") + self.assertFalse(result) + + def test_pipeline_does_not_deploy_when_gates_fail(self): + cfg = _make_config({"min_annotated_samples": 10}) + db = _make_db(annotated_count=50) + pipeline = ActiveLearningPipeline(db, cfg) + + with patch.object(pipeline._ingestor, "ingest", return_value="/data/training/v2.0.0"), \ + patch("active_learning.pipeline.Trainer") as MockTrainer, \ + patch("active_learning.pipeline.Validator") as MockValidator: + MockTrainer.return_value.train.return_value = "/models/candidate/v2.0.0_best.pt" + MockValidator.return_value.validate.return_value = { + "f1": 0.71, "ece": 0.05, + "gates_passed": False, + "gate_details": {"delta_f1": 0.01}, + } + result = pipeline.run("v2.0.0") + self.assertFalse(result) + + def test_pipeline_deploys_when_gates_pass(self): + cfg = _make_config({"min_annotated_samples": 10}) + db = _make_db(annotated_count=50) + pipeline = ActiveLearningPipeline(db, cfg) + + with patch.object(pipeline._ingestor, "ingest", return_value="/data/training/v2.0.0"), \ + patch("active_learning.pipeline.Trainer") as MockTrainer, \ + patch("active_learning.pipeline.Validator") as MockValidator, \ + patch.object(pipeline._registry, "deploy", return_value=True) as mock_deploy: + MockTrainer.return_value.train.return_value = "/models/candidate/v2.0.0_best.pt" + MockValidator.return_value.validate.return_value = { + "f1": 0.88, "ece": 0.04, + "gates_passed": True, + "gate_details": {"delta_f1": 0.18}, + } + result = pipeline.run("v2.0.0") + self.assertTrue(result) + mock_deploy.assert_called_once_with("v2.0.0", "/models/candidate/v2.0.0_best.pt") + + +if __name__ == "__main__": + unittest.main()