From 9f098d4b1ff26eaf30850cdb205325e61cfe3a69 Mon Sep 17 00:00:00 2001 From: Ryan Shpeherd Date: Wed, 9 Sep 2026 13:50:19 -0400 Subject: [PATCH] Update worker --- api/app.pl | 14 ++ src/api_client.py | 136 +++++++++++++++++ src/task_processor.py | 297 ++++++++++++++++++++++++++++++++++++++ src/task_worker.py | 186 ++++++++++++++++++++++++ worker/Dockerfile | 4 +- worker/docker-compose.yml | 40 +++++ 6 files changed, 675 insertions(+), 2 deletions(-) create mode 100644 src/api_client.py create mode 100644 src/task_processor.py create mode 100644 src/task_worker.py create mode 100644 worker/docker-compose.yml diff --git a/api/app.pl b/api/app.pl index e469589..67403fc 100644 --- a/api/app.pl +++ b/api/app.pl @@ -169,4 +169,18 @@ post '/api/v1/task/:task/complete' => sub { return { message => "Task completed successfully" }; }; +post '/api/v1/task' => sub { + my $video_id = body_parameters->get("video_id"); + my $task_type = body_parameters->get("task_type"); + + database->do("DELETE FROM tasks WHERE video_id=? AND task_type=?", undef, $video_id, $task_type); + + my $sth = database->prepare("INSERT INTO tasks (video_id, task_type, status) VALUES (?, ?, 'PENDING')"); + $sth->execute($video_id, $task_type); + my $task_id = database->last_insert_id(undef, undef, 'tasks', undef); + $sth->finish(); + + return { id => $task_id }; +}; + start(); \ No newline at end of file diff --git a/src/api_client.py b/src/api_client.py new file mode 100644 index 0000000..05c2608 --- /dev/null +++ b/src/api_client.py @@ -0,0 +1,136 @@ +"""HTTP client for VideoDetect's Perl REST API (api/app.pl). + +Wraps the three worker-facing endpoints: + GET /api/v1/nexttask/:type — claim next pending task + GET /api/v1/video/:id — fetch video metadata + POST /api/v1/task/:task/complete — submit results + +Uses the ``requests`` library; raises :class:`ApiError` on unexpected HTTP status codes. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass, field +from typing import Any, Optional + +import requests + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Exception +# --------------------------------------------------------------------------- + +class ApiError(Exception): + """Raised when the API returns an unexpected error status.""" + + def __init__(self, status_code: int, message: str): + self.status_code = status_code + self.message = message + super().__init__(f"API error {status_code}: {message}") + + +# --------------------------------------------------------------------------- +# Client +# --------------------------------------------------------------------------- + +class ApiClient: + """Thin HTTP client for the Dancer2 REST API.""" + + def __init__(self, base_url: str = "http://localhost:3000"): + self.base_url = base_url.rstrip("/") + self.session = requests.Session() + logger.info("API client configured: base_url=%s", self.base_url) + + # -- helpers ------------------------------------------------------------- + + def _get(self, path: str, **kwargs: Any) -> dict: + url = f"{self.base_url}{path}" + resp = self.session.get(url, **kwargs) + if resp.status_code == 404: + logger.debug("GET %s → 404", url) + return {} # caller distinguishes "not found" from real data + resp.raise_for_status() + return resp.json() + + def _post(self, path: str, json_body: dict[str, Any]) -> dict: + """POST a JSON body (Content-Type: application/json).""" + url = f"{self.base_url}{path}" + logger.debug("POST %s → %s", url, json.dumps(json_body)) + resp = self.session.post(url, json=json_body) + resp.raise_for_status() + return resp.json() + + # -- public API ---------------------------------------------------------- + + def get_next_task(self, task_type: str) -> Optional[dict[str, Any]]: + """Claim the next pending task. + + Returns ``{"task": {...}, "assign_key": "worker_NNN"}`` on success, + or ``None`` when no PENDING tasks remain (HTTP 404). + """ + path = f"/api/v1/nexttask/{task_type}" + result = self._get(path) + if not result: + return None + assert "task" in result and "assign_key" in result, \ + f"Unexpected response shape: {result}" + return result + + def get_video(self, video_id: int) -> dict[str, Any]: + """Fetch full video metadata by integer ID. + + Raises :class:`ApiError` if the video is not found (404). + """ + path = f"/api/v1/video/{video_id}" + result = self._get(path) + if not result: + raise ApiError(404, f"Video {video_id} not found") + assert "id" in result, f"Unexpected video response shape: {result}" + return result + + def submit_results( + self, + task_id: int, + assign_key: str, + results: dict[str, Any], + ) -> dict[str, Any]: + """Submit processing results for a claimed task. + + Args: + task_id: The database ID returned by ``get_next_task``. + assign_key: The worker token returned by ``get_next_task``. + results: A JSON-serializable dict (status, confidence, routing_decision, etc.). + + Returns the server response dict on success. + Raises :class:`ApiError` on 4xx/5xx. + """ + path = f"/api/v1/task/{task_id}/complete" + return self._post(path, { + "assign_key": assign_key, + "results": results, + }) + + def create_task(self, task_type: str, video_id: int) -> dict[str, Any]: + """Create a new pending task. + + Args: + task_type: Task type string (e.g. 'AISCAN', 'REVIEW'). + video_id: The video to associate the task with. + + Returns the server response dict on success. + Raises :class:`ApiError` on 4xx/5xx. + """ + path = "/api/v1/tasks" + return self._post(path, { + "task_type": task_type, + "video_id": video_id, + "status": "PENDING", + }) + + def close(self) -> None: + """Close the underlying HTTP session.""" + self.session.close() diff --git a/src/task_processor.py b/src/task_processor.py new file mode 100644 index 0000000..5de84a0 --- /dev/null +++ b/src/task_processor.py @@ -0,0 +1,297 @@ +"""Standalone AI-processing pipeline for the task-worker architecture. + +``process_video(video_dict, config)`` runs the full scan pipeline on a single +video file and returns a results dict that can be submitted via the REST API. + +It reuses existing classes from the codebase: + * ``VideoProber`` – ffprobe metadata extraction + * ``ScratchManager`` – per-video temp dirs under /scratch + * ``FrameSampler`` – uniform frame extraction via ffmpeg + * ``FaceDetector`` – TensorRT face detection on sampled frames + * ``FaceClassifier`` – MobileNetV3 classification of face crops + * ``aggregator.aggregate`` – per-crop → video-level confidence + * ``router.route`` – routing thresholds (MATCH / REVIEW / SKIP) + +GPU model loading is lazily initialised and cached at module level so that +sequential invocations within the same Python process reuse the same loaded +engines. This is important because each short-lived worker may be asked to +process several tasks in a row (``--tasks N``). +""" + +from __future__ import annotations + +import logging +import os +import time +from pathlib import Path +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Lazy imports — avoid heavy imports until actually needed +# --------------------------------------------------------------------------- +_face_detector_instance: Optional[Any] = None +_classifier_instance: Optional[Any] = None + + +def _get_face_detector(config: dict) -> Any: + """Return a singleton ``FaceDetector`` (initialised once per process).""" + global _face_detector_instance + if _face_detector_instance is not None: + return _face_detector_instance + + from face_detector import FaceDetector # noqa: local import, heavy + + cfg = (config or {}).get("face_detection", {}) + engine_path = cfg.get( + "model_path", "/models/face_detector/face_detector.trt" + ) + _face_detector_instance = FaceDetector( + engine_path=engine_path, + input_size=int(cfg.get("input_size", 640)), + confidence_threshold=float(cfg.get("confidence_threshold", 0.25)), + iou_threshold=float(cfg.get("iou_threshold", 0.45)), + max_faces_per_frame=int(cfg.get("max_faces_per_frame", 10)), + max_faces_per_video=int(cfg.get("max_faces_per_video", 100)), + ) + return _face_detector_instance + + +def _get_classifier(config: dict) -> Any: + """Return a singleton ``FaceClassifier`` (initialised once per process).""" + global _classifier_instance + if _classifier_instance is not None: + return _classifier_instance + + from classifier import FaceClassifier # noqa: local import, heavy + + cfg = (config or {}).get("classifier", {}) + engine_path = cfg.get( + "model_path", "/models/classifier/classifier.trt" + ) + _classifier_instance = FaceClassifier( + engine_path=engine_path, + temperature=float(cfg.get("temperature", 1.0)), + input_size=int(cfg.get("input_size", 224)), + ) + return _classifier_instance + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def get_model_version() -> str: + """Return the currently active model version from the DB (best-effort).""" + try: + from db_connector import DatabaseConnector # noqa: local import, heavy + db = DatabaseConnector("default") + row = 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 process_video(video: Dict[str, Any], config: Optional[Dict] = None) -> Dict[str, Any]: + """Run the full AI scan pipeline on a single video file. + + Args: + video: Metadata dict (as returned by ``ApiClient.get_video``). + Must contain at minimum ``"id"`` and ``"file_path"``. + config: The loaded ``config.yaml`` dict (or None for defaults). + + Returns: + A results dict suitable for submission via the REST API. + """ + config = config or {} + start_time = time.time() + video_id = video["id"] + file_path = video["file_path"] + + logger.info("Processing video %d: %s", video_id, file_path) + + sampling_cfg = (config or {}).get("sampling", {}) + storage_cfg = (config or {}).get("storage", {}) + gpu_cfg = (config or {}).get("gpu", {}) + batching_cfg = (config or {}).get("batching", {}) + agg_cfg = (config or {}).get("aggregation", {}) + routing_cfg = (config or {}).get("routing", {}) + + try: + # 1. Probe metadata (use DB metadata if already populated, re-probe) + prober = VideoProber(timeout=10) + metadata = prober.probe(file_path) + if metadata.is_unscannable: + raise RuntimeError(metadata.error_message or "Video metadata could not be determined") + + # Use video from API if available, fall back to probe results + resolution_w = video.get("resolution_w") or metadata.resolution_w or 640 + resolution_h = video.get("resolution_h") or metadata.resolution_h or 480 + duration = video.get("duration") or metadata.duration or 10.0 + + logger.info( + "Video %d: codec=%s resolution=%dx%d duration=%.1fs", + video_id, metadata.codec or "unknown", + resolution_w, resolution_h, duration, + ) + + # 2. Set up scratch space + scratch = ScratchManager( + base_path=storage_cfg.get("scratch_path", "/scratch"), + video_id=str(video_id), + auto_cleanup=True, + ) + frame_dir = scratch.ensure_frame_dir() + + # 3. Sample frames + interval = int(sampling_cfg.get("interval_seconds", 30)) + quality = int(sampling_cfg.get("quality", 2)) + sampler = FrameSampler(interval_seconds=interval, quality=quality) + + resolution = (resolution_w, resolution_h) + extracted_frames = sampler.extract_frames( + video_path=file_path, + output_dir=str(frame_dir), + duration=duration, + interval_seconds=interval, + resolution=resolution, + timeout_seconds=30, + ) + + if not extracted_frames: + scratch.cleanup_all() + raise RuntimeError("No frames were extracted") + + logger.info( + "Video %d: extracted %d frame(s)", video_id, len(extracted_frames) + ) + + # 4. Face detection + detector = _get_face_detector(config) + batch_size = _get_batch_size(gpu_cfg, batching_cfg) + + detections_per_frame = detector.detect_faces( + extracted_frames, batch_size=batch_size + ) + + all_detections = sorted( + [d for dets in detections_per_frame for d in dets], + key=lambda d: d.confidence, + reverse=True, + ) + total_faces = len(all_detections) + + # 5. Classify (or SKIP if no faces) + if total_faces == 0: + logger.info( + "Video %d: no faces detected → routing to SKIP", video_id + ) + scratch.cleanup_all() + return _failed_result(video_id, start_time, 0.0, "SKIP") + + crop_dir = Path(scratch.frame_dir.parent) / "crops" + cropped_detections = detector.extract_crops( + all_detections, output_dir=str(crop_dir), crop_size=(224, 224) + ) + crop_paths = [d.crop_path for d in cropped_detections if d.crop_path] + + classifier = _get_classifier(config) + frame_confidences = classifier.classify( + crop_paths, batch_size=batch_size + ) + + # 6. Aggregate confidence + from aggregator import aggregate as agg_func # noqa: local import + + video_confidence = agg_func( + frame_confidences, + strategy=agg_cfg.get("strategy", "max"), + alpha=float(agg_cfg.get("alpha", 1.0)), + beta=float(agg_cfg.get("beta", 0.1)), + top_k=int(agg_cfg.get("top_k", 3)), + ) + + # 7. Route + from router import route as route_fn # noqa: local import + + routing = route_fn( + video_confidence, + t_high=float(routing_cfg.get("T_high", 0.75)), + t_low=float(routing_cfg.get("T_low", 0.45)), + ) + + processing_time = time.time() - start_time + logger.info( + "Video %d complete: C=%.4f routing=%s faces=%d frames=%d time=%.1fs", + video_id, video_confidence, routing, total_faces, + len(extracted_frames), processing_time, + ) + + scratch.cleanup_all() + + return { + "status": "COMPLETED", + "confidence": round(video_confidence, 4), + "routing_decision": routing, + "face_count": total_faces, + "frame_count": len(extracted_frames), + "model_version": get_model_version(), + "processing_time_seconds": round(processing_time, 2), + "error": None, + } + + except Exception as exc: + logger.error("Processing video %d failed: %s", video_id, exc, exc_info=True) + try: + scratch.cleanup_all() # noqa: undefined-name guard below + except NameError: + pass # scratch was never created (probe/probe error) + + return { + "status": "FAILED", + "confidence": 0.0, + "routing_decision": "REVIEW", + "face_count": 0, + "frame_count": 0, + "model_version": get_model_version(), + "processing_time_seconds": round(time.time() - start_time, 2), + "error": str(exc), + } + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _get_batch_size(gpu_cfg: dict, batching_cfg: dict) -> int: + """Compute an appropriate batch size based on GPU config.""" + from gpu_manager import GPUMemoryManager # noqa: local import, heavy + + mgr = GPUMemoryManager( + max_memory_gb=gpu_cfg.get("max_memory_gb", 18.0), + reduce_threshold_gb=batching_cfg.get("vram_reduce_threshold_gb", 16.0), + increase_threshold_gb=batching_cfg.get("vram_increase_threshold_gb", 10.0), + initial_batch_size=batching_cfg.get("max_batch_size", 16), + ) + return mgr.current_batch_size + + +def _failed_result( + video_id: int, start_time: float, confidence: float, routing: str +) -> Dict[str, Any]: + """Build a SKIP/early-exit result dict (no faces case).""" + return { + "status": "COMPLETED", + "confidence": round(confidence, 4), + "routing_decision": routing, + "face_count": 0, + "frame_count": 0, + "model_version": get_model_version(), + "processing_time_seconds": round(time.time() - start_time, 2), + "error": None, + } diff --git a/src/task_worker.py b/src/task_worker.py new file mode 100644 index 0000000..317dcfc --- /dev/null +++ b/src/task_worker.py @@ -0,0 +1,186 @@ +"""VideoDetect task worker — short-lived process for AI scanning. + +Usage:: + + python3 -m src.task_worker # process 1 task and exit + python3 -m src.task_worker --tasks 50 # process up to 50 tasks + python3 src/task_worker.py --tasks 10 # direct invocation + +Environment variables: + API_BASE_URL (default http://localhost:3000) + Base URL of the Dancer2 REST API. + + TASK_COUNT (default 1) + Maximum number of tasks to process before exiting. + +On success exits with code 0; on unrecoverable errors exits non-zero. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import signal +import sys +from pathlib import Path + +# Ensure the src/ directory is on sys.path so our submodules import cleanly +_src = Path(__file__).resolve().parent +if str(_src) not in sys.path: + sys.path.insert(0, str(_src)) + +from api_client import ApiClient, ApiError +from task_processor import process_video, get_model_version +from config_loader import get_config +from logging_config import setup_logging + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Graceful shutdown +# --------------------------------------------------------------------------- + +_shutdown_requested = False + + +def _signal_handler(signum: int, frame) -> None: + logger.info("Received signal %d — finishing current task then exiting", signum) + global _shutdown_requested + _shutdown_requested = True + + +signal.signal(signal.SIGTERM, _signal_handler) +signal.signal(signal.SIGINT, _signal_handler) + +# --------------------------------------------------------------------------- +# CLI entry point +# --------------------------------------------------------------------------- + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="VideoDetect task worker — process AI scan tasks", + ) + parser.add_argument( + "--tasks", + type=int, + default=int(os.environ.get("TASK_COUNT", "1")), + help="Number of tasks to process before exiting (default: 1 or $TASK_COUNT)", + ) + parser.add_argument( + "--api-url", + default=os.environ.get("API_BASE_URL", "http://localhost:3000"), + help="Base URL of the Dancer2 REST API (default: http://localhost:3000)", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + """Run the task worker loop. + + Returns exit code (0 = success). + """ + args = parse_args(argv) + + # Load config for the processing pipeline + config = get_config() + + # Setup JSON logging + log_cfg = (config or {}).get("logging", {}) + setup_logging( + level=log_cfg.get("level", "INFO"), + log_format=log_cfg.get("format", "json"), + rotation_max_bytes=log_cfg.get("rotation_max_bytes", 104857600), + rotation_backup_count=log_cfg.get("rotation_backup_count", 10), + ) + + model_version = get_model_version() + logger.info( + "Task worker starting. Will process up to %d task(s). Model version: %s", + args.tasks, + model_version, + ) + + api = ApiClient(base_url=args.api_url) + total_attempted = 0 + total_succeeded = 0 + total_failed = 0 + + try: + for i in range(1, args.tasks + 1): + if _shutdown_requested: + logger.info("Shutdown requested after %d/%d tasks", i - 1, args.tasks) + break + + logger.info("=== Processing task %d/%d ===", i, args.tasks) + + # 1. Claim a task + response = api.get_next_task("AISCAN") + if response is None: + logger.info("No more tasks available (API returned no task)") + break + + task_id = response["task"]["id"] + assign_key = response["assign_key"] + video_id = response["task"]["video_id"] + total_attempted += 1 + + logger.info( + "Claiming task %d (video_id=%d) for AISCAN", task_id, video_id + ) + + # 2. Fetch video metadata + try: + video = api.get_video(video_id) + except ApiError as exc: + logger.error("Failed to fetch video %d via API: %s", video_id, exc) + total_failed += 1 + continue + + # 3. Process the video (AI pipeline) + results = process_video(video, config=config) + + if results["status"] == "COMPLETED": + logger.info( + "Video %d complete: C=%.4f routing=%s faces=%d frames=%d time=%.1fs", + video_id, + results.get("confidence", 0), + results.get("routing_decision", "?"), + results.get("face_count", 0), + results.get("frame_count", 0), + results.get("processing_time_seconds", 0), + ) + total_succeeded += 1 + else: + logger.warning( + "Video %d failed: %s", + video_id, + results.get("error", "unknown error"), + ) + total_failed += 1 + + # 4. Submit results via API + try: + api.submit_results(task_id, assign_key, results) + logger.info("Submitted results for task %d via API", task_id) + except ApiError as exc: + logger.error( + "Failed to submit results for task %d: %s", task_id, exc + ) + + finally: + api.close() + + logger.info( + "Task worker finished: %d attempted, %d succeeded, %d failed", + total_attempted, + total_succeeded, + total_failed, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/worker/Dockerfile b/worker/Dockerfile index 5f35ff6..5bb7a2f 100644 --- a/worker/Dockerfile +++ b/worker/Dockerfile @@ -64,5 +64,5 @@ HEALTHCHECK --interval=30s --timeout=10s --retries=3 \ # Switch to non-root user USER appuser -# Default command -CMD ["python3", "-m", "src.main"] +# Default command: process tasks from the API until drained +CMD ["python3", "-m", "src.task_worker"] diff --git a/worker/docker-compose.yml b/worker/docker-compose.yml new file mode 100644 index 0000000..1e894ac --- /dev/null +++ b/worker/docker-compose.yml @@ -0,0 +1,40 @@ +# VideoDetect — Worker service (standalone compose file) +# Run: docker compose -f worker/docker-compose.yml up -d + +services: + worker: + build: + context: ./ + dockerfile: Dockerfile + container_name: videodetect-worker + restart: "no" + environment: + - API_BASE_URL=${API_BASE_URL:-http://videodetect-api:3000} + - TASK_COUNT=${TASK_COUNT:-50} + - REVIEW_CONFIDENCE=${REVIEW_CONFIDENCE:-0.75} + - LOG_FILE=/logs/videodetect.log + volumes: + - nas_input:/data/input:ro + - scratch_data:/scratch + - $PWD/worker-logs:/logs + networks: + - videodetect_videodetect-network + deploy: + resources: + limits: + memory: 4G + nvidia-gpus: "1" + +volumes: + nas_input: + driver: local + driver_opts: + type: nfs + o: addr=10.0.0.2,ro,nfsvers=4,hard,intr + device: ":/mnt/Bulk/Homes/ryan/Prawns" + scratch_data: + driver: local + +networks: + videodetect_videodetect-network: + external: true