Story 6
This commit is contained in:
Ryan Shpeherd
@@ -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
|
||||
+44
-3
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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():
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user