Ryan/VideoDetect
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update worker

Browse Source
This commit is contained in:
Ryan Shpeherd
2026-09-09 13:50:19 -04:00
parent d9c0998400
commit 9f098d4b1f
6 changed files with 675 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/app.pl
+14
View File
@@ -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();
src/api_client.py
+136
View File
@@ -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()
src/task_processor.py
+297
View File
@@ -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,
}
src/task_worker.py
+186
View File
@@ -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())
worker/Dockerfile
+2 -2
View File
@@ -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"]
worker/docker-compose.yml
+40
View File
@@ -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