diff --git a/src/data_export.py b/src/data_export.py new file mode 100644 index 0000000..474cf7c --- /dev/null +++ b/src/data_export.py @@ -0,0 +1,115 @@ +"""Batch Parquet/JSONL export of video processing results.""" + +import json +import logging +import os +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + + +class DataExporter: + """Buffer result records and flush to Parquet or JSONL when the batch is full.""" + + def __init__( + self, + output_path: str, + model_version: str = "v0.0.0", + export_format: str = "parquet", + compression: str = "snappy", + batch_size: int = 100, + include_frame_confidences: bool = True, + ): + self.output_path = Path(output_path) / model_version + self.model_version = model_version + self.export_format = export_format.lower() + self.compression = compression + self.batch_size = batch_size + self.include_frame_confidences = include_frame_confidences + self._buffer: List[Dict[str, Any]] = [] + self.output_path.mkdir(parents=True, exist_ok=True) + + def add(self, record: Dict[str, Any]) -> None: + """Buffer one result record; flush automatically when batch is full.""" + self._buffer.append(record) + if len(self._buffer) >= self.batch_size: + self.flush() + + def flush(self) -> Optional[str]: + """Write buffered records to disk; returns the output file path or None.""" + if not self._buffer: + return None + + records = self._buffer[:] + self._buffer.clear() + + if not self.include_frame_confidences: + for r in records: + r.pop("confidence_scores", None) + + batch_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + stem = f"{self.model_version}_{batch_id}" + + if self.export_format in ("parquet", "both"): + path = self._write_parquet(records, stem) + if self.export_format in ("jsonl", "both"): + path = self._write_jsonl(records, stem) + if self.export_format not in ("parquet", "jsonl", "both"): + logger.warning("Unknown export format '%s'; defaulting to jsonl", self.export_format) + path = self._write_jsonl(records, stem) + + return str(path) + + def _write_parquet(self, records: List[Dict[str, Any]], stem: str) -> Path: + out = self.output_path / f"{stem}.parquet" + try: + import pyarrow as pa + import pyarrow.parquet as pq + + schema = pa.schema([ + pa.field("video_id", pa.int64()), + pa.field("file_path", pa.string()), + pa.field("model_version", pa.string()), + pa.field("sample_count", pa.int32()), + pa.field("confidence_scores", pa.list_(pa.float64())), + pa.field("video_confidence", pa.float64()), + pa.field("routing", pa.string()), + pa.field("processed_at", pa.string()), + ]) + + table = pa.table( + { + "video_id": [r.get("video_id") for r in records], + "file_path": [r.get("file_path", "") for r in records], + "model_version": [r.get("model_version", self.model_version) for r in records], + "sample_count": [r.get("sample_count", 0) for r in records], + "confidence_scores": [r.get("confidence_scores", []) for r in records], + "video_confidence": [float(r.get("video_confidence", 0.0)) for r in records], + "routing": [r.get("routing", "SKIP") for r in records], + "processed_at": [r.get("processed_at", "") for r in records], + }, + schema=schema, + ) + pq.write_table(table, out, compression=self.compression) + logger.info("Exported %d records to %s", len(records), out) + except ImportError: + logger.warning("pyarrow not available; falling back to JSONL") + out = self._write_jsonl(records, stem.replace(".parquet", "")) + return out + + def _write_jsonl(self, records: List[Dict[str, Any]], stem: str) -> Path: + out = self.output_path / f"{stem}.jsonl" + with open(out, "w", encoding="utf-8") as fh: + for record in records: + fh.write(json.dumps(record, default=str) + "\n") + logger.info("Exported %d records to %s", len(records), out) + return out + + def __del__(self): + if self._buffer: + try: + self.flush() + except Exception: + pass diff --git a/src/orchestrator.py b/src/orchestrator.py index dcbecae..479549c 100644 --- a/src/orchestrator.py +++ b/src/orchestrator.py @@ -15,10 +15,12 @@ from typing import Dict, List, Optional from batcher import DynamicBatcher from classifier import FaceClassifier +from data_export import DataExporter from face_detector import Detection, FaceDetector from frame_sampler import FrameSampler from gpu_manager import GPUMemoryManager from prober import VideoProber +from result_updater import ResultUpdater from scratch_manager import ScratchManager import aggregator import router @@ -111,6 +113,27 @@ class WorkerPool: ) self._aggregation_config = (config or {}).get("aggregation", {}) self._routing_config = (config or {}).get("routing", {}) + model_version = self._get_active_model_version() + export_config = (config or {}).get("export", {}) + self._result_updater = ResultUpdater(db_connector, model_version=model_version) + self._exporter = DataExporter( + output_path=export_config.get("output_path", "/data/output"), + model_version=model_version, + export_format=export_config.get("format", "parquet"), + compression=export_config.get("compression", "snappy"), + batch_size=int(export_config.get("batch_size", 100)), + include_frame_confidences=bool(export_config.get("include_frame_confidences", True)), + ) + + def _get_active_model_version(self) -> str: + """Return the currently active model version from the DB.""" + try: + row = self.db.fetchone("SELECT version FROM models WHERE status = 'ACTIVE' LIMIT 1") + if row: + return row["version"] + except Exception as exc: + logger.debug("Could not fetch active model version: %s", exc) + return "v0.0.0-placeholder" def start(self): """Start the worker pool.""" @@ -267,13 +290,31 @@ class WorkerPool: len(cropped_detections), len(extracted_frames), ) - self._complete_job( - job, + # Persist results atomically (video update + processing log) before cleanup + persisted = self._result_updater.persist( + video_id=job.video_id, frame_count=len(extracted_frames), confidence=video_confidence, routing=routing_decision, + frame_confidences=frame_confidences, ) - scratch_manager.cleanup() + + if persisted: + job.status = JobStatus.COMPLETED + job.completed_at = datetime.now(timezone.utc) + self._exporter.add({ + "video_id": job.video_id, + "file_path": job.file_path, + "model_version": self._result_updater.model_version, + "sample_count": len(extracted_frames), + "confidence_scores": frame_confidences, + "video_confidence": video_confidence, + "routing": routing_decision, + "processed_at": datetime.now(timezone.utc).isoformat(), + }) + + # Scratch cleanup only after successful persistence + scratch_manager.cleanup_all() self._jobs_processed += 1 return True diff --git a/src/processing_logger.py b/src/processing_logger.py new file mode 100644 index 0000000..75499af --- /dev/null +++ b/src/processing_logger.py @@ -0,0 +1,39 @@ +"""Insert processing audit log entries within an existing DB transaction.""" + +import json +import logging +from datetime import datetime, timezone +from typing import List, Optional + +logger = logging.getLogger(__name__) + + +def insert_log( + cursor, + video_id: int, + model_version: str, + frame_count: int, + confidence_score: Optional[float], + routing_decision: str, + frame_confidences: Optional[List[float]] = None, +) -> None: + """Insert one row into processing_logs; must be called inside an open transaction.""" + now = datetime.now(timezone.utc).replace(tzinfo=None) + confidence_scores_json = json.dumps(frame_confidences) if frame_confidences is not None else None + + cursor.execute( + """INSERT INTO processing_logs + (video_id, model_version, frame_count, confidence_score, + confidence_scores, routing_decision, processed_at) + VALUES (%s, %s, %s, %s, %s, %s, %s)""", + ( + video_id, + model_version, + frame_count or 0, + confidence_score, + confidence_scores_json, + routing_decision, + now, + ), + ) + logger.debug("Inserted processing log for video %s (routing=%s)", video_id, routing_decision) diff --git a/src/result_updater.py b/src/result_updater.py new file mode 100644 index 0000000..3dd1ffb --- /dev/null +++ b/src/result_updater.py @@ -0,0 +1,75 @@ +"""Atomic result persistence: update videos + insert processing_logs in one transaction.""" + +import logging +from datetime import datetime, timezone +from typing import List, Optional + +import processing_logger + +logger = logging.getLogger(__name__) + + +class ResultUpdater: + def __init__(self, db_connector, model_version: str = "v0.0.0-placeholder"): + self.db = db_connector + self.model_version = model_version + + def persist( + self, + video_id: int, + frame_count: int, + confidence: float, + routing: str, + frame_confidences: Optional[List[float]] = None, + ) -> bool: + """Atomically update videos and insert a processing log. + + Returns False without raising if the state guard prevents the update + (video is no longer in PROCESSING state). + """ + now = datetime.now(timezone.utc).replace(tzinfo=None) + + try: + with self.db.transaction() as conn: + cursor = conn.cursor() + try: + cursor.execute( + """UPDATE videos + SET status = 'COMPLETED', + last_processed_time = %s, + confidence_score = %s, + routing_decision = %s, + model_version = %s, + frame_count = %s, + updated_at = %s + WHERE id = %s AND status = 'PROCESSING'""", + (now, confidence, routing, self.model_version, frame_count, now, video_id), + ) + + if cursor.rowcount == 0: + logger.warning( + "State guard: video %s is not PROCESSING; skipping update", video_id + ) + return False + + processing_logger.insert_log( + cursor, + video_id=video_id, + model_version=self.model_version, + frame_count=frame_count, + confidence_score=confidence, + routing_decision=routing, + frame_confidences=frame_confidences, + ) + finally: + cursor.close() + + logger.info( + "Persisted results for video %s: C=%.4f routing=%s model=%s", + video_id, confidence, routing, self.model_version, + ) + return True + + except Exception as exc: + logger.error("Failed to persist results for video %s: %s", video_id, exc) + raise diff --git a/src/scratch_manager.py b/src/scratch_manager.py index 7d6fc1f..a58d4b2 100644 --- a/src/scratch_manager.py +++ b/src/scratch_manager.py @@ -23,11 +23,18 @@ class ScratchManager: return self.frame_dir def cleanup(self) -> None: - """Remove the video-specific frame directory and everything underneath it.""" + """Remove the frame scratch directory only.""" if self.auto_cleanup and self.frame_dir.exists(): shutil.rmtree(self.frame_dir) logger.info("Cleaned scratch frames for %s", self.video_id) + def cleanup_all(self) -> None: + """Remove the entire per-video scratch directory (frames + crops).""" + video_root = self.base_path / self.video_id + if self.auto_cleanup and video_root.exists(): + shutil.rmtree(video_root) + logger.info("Cleaned all scratch data for %s", self.video_id) + def usage_bytes(self) -> int: """Return the total size of the scratch directory in bytes.""" if not self.frame_dir.exists(): diff --git a/tests/test_story_06.py b/tests/test_story_06.py new file mode 100644 index 0000000..5f9e4b3 --- /dev/null +++ b/tests/test_story_06.py @@ -0,0 +1,148 @@ +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +import processing_logger +from data_export import DataExporter +from result_updater import ResultUpdater +from scratch_manager import ScratchManager + + +class Story06ProcessingLoggerTests(unittest.TestCase): + def test_insert_log_executes_correct_sql(self): + cursor = MagicMock() + processing_logger.insert_log( + cursor, + video_id=42, + model_version="v1.0.0", + frame_count=8, + confidence_score=0.82, + routing_decision="MATCH", + frame_confidences=[0.80, 0.82, 0.85], + ) + cursor.execute.assert_called_once() + sql, params = cursor.execute.call_args[0] + self.assertIn("INSERT INTO processing_logs", sql) + self.assertEqual(params[0], 42) # video_id + self.assertEqual(params[1], "v1.0.0") # model_version + self.assertEqual(params[2], 8) # frame_count + self.assertAlmostEqual(params[3], 0.82) # confidence_score + scores = json.loads(params[4]) # confidence_scores JSON + self.assertEqual(scores, [0.80, 0.82, 0.85]) + self.assertEqual(params[5], "MATCH") # routing_decision + + def test_insert_log_null_frame_confidences(self): + cursor = MagicMock() + processing_logger.insert_log(cursor, 1, "v0", 0, None, "SKIP") + _, params = cursor.execute.call_args[0] + self.assertIsNone(params[4]) # confidence_scores column + + +class Story06ResultUpdaterTests(unittest.TestCase): + def _make_db(self, rowcount=1): + cursor = MagicMock() + cursor.rowcount = rowcount + conn = MagicMock() + conn.cursor.return_value = cursor + db = MagicMock() + db.transaction.return_value.__enter__ = MagicMock(return_value=conn) + db.transaction.return_value.__exit__ = MagicMock(return_value=False) + return db, cursor + + def test_persist_returns_true_when_update_succeeds(self): + db, cursor = self._make_db(rowcount=1) + updater = ResultUpdater(db, model_version="v1.0.0") + result = updater.persist( + video_id=7, frame_count=5, confidence=0.9, + routing="MATCH", frame_confidences=[0.9], + ) + self.assertTrue(result) + + def test_persist_returns_false_on_state_guard_miss(self): + db, cursor = self._make_db(rowcount=0) + updater = ResultUpdater(db, model_version="v1.0.0") + result = updater.persist( + video_id=7, frame_count=5, confidence=0.9, + routing="MATCH", frame_confidences=[0.9], + ) + self.assertFalse(result) + + def test_persist_calls_insert_log_after_update(self): + db, cursor = self._make_db(rowcount=1) + updater = ResultUpdater(db, model_version="v1.0.0") + updater.persist(video_id=7, frame_count=5, confidence=0.9, + routing="MATCH", frame_confidences=[0.9]) + # cursor.execute called twice: UPDATE videos + INSERT processing_logs + self.assertEqual(cursor.execute.call_count, 2) + + +class Story06DataExporterTests(unittest.TestCase): + def test_jsonl_flush_writes_valid_records(self): + with tempfile.TemporaryDirectory() as tmpdir: + exporter = DataExporter( + output_path=tmpdir, model_version="v1", + export_format="jsonl", batch_size=100, + ) + exporter.add({"video_id": 1, "routing": "MATCH", "video_confidence": 0.9, + "confidence_scores": [0.9], "sample_count": 1, + "file_path": "/a.mp4", "processed_at": "2026-01-01T00:00:00+00:00"}) + path = exporter.flush() + + self.assertIsNotNone(path) + lines = Path(path).read_text(encoding="utf-8").strip().split("\n") + self.assertEqual(len(lines), 1) + record = json.loads(lines[0]) + self.assertEqual(record["routing"], "MATCH") + self.assertAlmostEqual(record["video_confidence"], 0.9) + + def test_auto_flush_at_batch_size(self): + with tempfile.TemporaryDirectory() as tmpdir: + exporter = DataExporter( + output_path=tmpdir, model_version="v1", + export_format="jsonl", batch_size=2, + ) + exporter.add({"video_id": 1, "routing": "SKIP", "video_confidence": 0.1, + "confidence_scores": [], "sample_count": 0, + "file_path": "/a.mp4", "processed_at": ""}) + exporter.add({"video_id": 2, "routing": "MATCH", "video_confidence": 0.9, + "confidence_scores": [0.9], "sample_count": 1, + "file_path": "/b.mp4", "processed_at": ""}) + # batch_size=2 → auto-flush triggered on second add + self.assertEqual(exporter._buffer, []) + + def test_exclude_frame_confidences_when_disabled(self): + with tempfile.TemporaryDirectory() as tmpdir: + exporter = DataExporter( + output_path=tmpdir, model_version="v1", + export_format="jsonl", batch_size=100, + include_frame_confidences=False, + ) + exporter.add({"video_id": 1, "routing": "SKIP", "video_confidence": 0.1, + "confidence_scores": [0.1, 0.2], "sample_count": 2, + "file_path": "/a.mp4", "processed_at": ""}) + path = exporter.flush() + record = json.loads(Path(path).read_text()) + self.assertNotIn("confidence_scores", record) + + +class Story06ScratchManagerCleanupAllTests(unittest.TestCase): + def test_cleanup_all_removes_entire_video_directory(self): + with tempfile.TemporaryDirectory() as tmpdir: + manager = ScratchManager(base_path=tmpdir, video_id="v42", auto_cleanup=True) + frame_dir = manager.ensure_frame_dir() + crops_dir = Path(tmpdir) / "v42" / "crops" + crops_dir.mkdir(parents=True) + (frame_dir / "frame.jpg").write_bytes(b"f") + (crops_dir / "crop.jpg").write_bytes(b"c") + + manager.cleanup_all() + self.assertFalse((Path(tmpdir) / "v42").exists()) + + +if __name__ == "__main__": + unittest.main()