18 Commits
Author SHA1 Message Date
Ryan Shpeherd 5e5972c31a Fix scanner task creator 2026-09-26 10:59:52 -04:00
Ryan Shpeherd a04e6e241b Updates to ai worker 2026-09-26 10:23:55 -04:00
Ryan Shpeherd e443191833 Add model and path to results 2026-09-25 22:43:54 -04:00
Ryan Shpeherd 589b775cda Fixes and new ai model 2026-09-25 22:40:24 -04:00
Ryan Shpeherd 8f2834775b Fixes 2026-09-15 13:59:41 -04:00
Ryan Shpeherd 75f38ec7ed Fixes 2026-09-15 13:28:43 -04:00
Ryan Shpeherd 5fa67a733c Fixes 2026-09-15 13:20:48 -04:00
Ryan Shpeherd 46548881cb Fixes 2026-09-15 13:12:21 -04:00
Ryan Shpeherd 4782dc5bfa Fix path 2026-09-15 12:58:16 -04:00
Ryan Shpeherd 1611485a66 Local 2026-09-15 12:39:47 -04:00
Ryan Shpeherd 639c29c2c4 Fix scanner 2026-09-14 16:02:21 -04:00
Ryan Shpeherd 3eb5b5b08b Add scanner step to set review tasks 2026-09-14 14:42:54 -04:00
Ryan 4be2684830 Fix query 2026-09-13 19:07:25 -04:00
Ryan ad3deb3226 Remote network usage 2026-09-13 18:40:02 -04:00
Ryan dbbbaa2ea2 Fix config password 2026-09-13 15:17:31 -04:00
Ryan a9d3143a8e Fix variables 2026-09-13 15:15:20 -04:00
Ryan 01a5e50fcd AI Scanning 2026-09-13 14:38:14 -04:00
Ryan Shpeherd 458d8862cc Cleanup 2026-09-09 15:01:42 -04:00
61 changed files with 841 additions and 8023 deletions
+2 -2
View File
@@ -4,8 +4,8 @@
# ----------------------------------------------------- # -----------------------------------------------------
# Database # Database
# ----------------------------------------------------- # -----------------------------------------------------
DB_ROOT_PASSWORD=changeme_root DB_ROOT_PASSWORD=videodetect
DB_PASSWORD=changeme_videodetect DB_PASSWORD=videodetect
DB_USER=videodetect DB_USER=videodetect
DB_NAME=videodetect DB_NAME=videodetect
DB_HOST=mariadb DB_HOST=mariadb
+5
View File
@@ -0,0 +1,5 @@
# VideoDetect — Docker Compose environment variables
# Copy this file to .env and fill in real values before running docker-compose up.
API_HOST=https://10.0.1.34:3000
+29
View File
@@ -0,0 +1,29 @@
FROM ubuntu:22.04
# Avoid interactive prompts during build
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC
# System dependencies
# Install Perl, system deps, and ffmpeg (includes ffprobe)
RUN apt update && apt install -y \
libdbd-mysql-perl \
libhttp-lite-perl \
libjson-perl \
libyaml-perl \
perl \
libdbi-perl \
build-essential \
ffmpeg \
libstring-shellquote-perl \
libwww-perl \
liblwp-protocol-https-perl
RUN mkdir -p /app
ADD run.pl /app/run.pl
RUN chmod +x /app/run.pl
WORKDIR /app
# Default command: process tasks from the API until drained
CMD ["perl", "/app/run.pl"]
+27
View File
@@ -0,0 +1,27 @@
# VideoDetect — Worker service (standalone compose file)
# Run: docker compose -f aiworker/docker-compose.yml up -d
services:
worker:
build:
context: ./
dockerfile: Dockerfile
container_name: videodetect-aiworker
restart: "no"
environment:
- API_BASE_URL=${API_BASE_URL:-http://10.0.1.34:3000}
- OLLAMA_URL=${OLLAMA_URL:-http://10.0.1.7:11434}
- SECONDS_PER_FRAME=${SECONDS_PER_FRAME:-30}
- TASK_COUNT=${TASK_COUNT:-50}
volumes:
- nas_input:/data:ro
- ./run.pl:/app/run.pl:ro
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"
+178
View File
@@ -0,0 +1,178 @@
#!/usr/bin/perl
use strict;
use MIME::Base64 qw(encode_base64);
use LWP::UserAgent;
use JSON;
use Data::Dumper;
my $seconds = $ENV{'SECONDS_PER_FRAME'} || 30;
my $ollama = $ENV{'OLLAMA_URL'} || 'http://10.0.1.7:11434';
my $api = $ENV{'API_BASE_URL'} || 'http://10.0.1.7:3000';
my $task_count = $ENV{'TASK_COUNT'} || 1;
my $image_count = $ENV{'IMAGE_COUNT'} || 10; # Number of images to extract from the video
my $min_time = $ENV{'MIN_TIME'} || 5; # Minimum time between images
my $sensitivity = 20;
my $resolution = "672:672";
#my $model = "hf.co/Jommarn/UNSEEN_Gemma_4_26B_NSFW-GGUF:Q8_0";
my $model = "hf.co/Jommarn/UNSEEN_Gemma_4_26B_NSFW-GGUF:Q5_K_M";
chdir("/tmp") || die "Unable to change directory to /tmp: $!\n";
my $ua = LWP::UserAgent->new;
$ua->timeout(60);
if(@ARGV) {
CheckPath($ARGV[0], {'duration' => 0});
exit;
}
for(my $i=0; $i<$task_count; $i++) {
my $ref = GetNext();
my $task = $ref->{'task'};
my $video = $ref->{'video'};
my $task_id = $task->{'id'};
my $video_id = $video->{'id'};
my $path = $video->{'file_path'};
my $result = CheckPath($path, $video);
next unless $result;
my $response = PostResponse($task_id, $ref->{'assign_key'}, $result);
}
print "Done\n";
exit(0);
sub CheckPath {
my ($path, $video) = @_;
die "File path not set\n" unless $path;
die "File does not exist: $path\n" unless -e $path;
unlink(<*.png>);
# Pick a frame interval that spreads $image_count frames across the video,
# but never more frequently than 1 frame per $min_time seconds.
my $interval = $seconds; # fallback if duration is unknown
if (my $duration = $video->{'duration'}) {
my $desired = $duration / $image_count;
$interval = ($desired > $min_time) ? $desired : $min_time;
}
system("ffmpeg -i \"$path\" -vf fps=1/$interval,scale=$resolution output_%04d.png 2>&1");
my $total=0;
my $yes=0;
$/ = undef;
foreach my $file (<*.png>) {
print "$file...\n";
open(IN, '<:raw', $file) || die "Unable to read $file: $!\n";
my $data = <IN>;
close(IN);
my $base64 = encode_base64($data);
$data=undef;
my $ref = {
'model' => $model,
'prompt' => "Are there any african american men in this photo, whether they are shown completely or not? Judge based on skin tone, accounting for shadows. Answer yes for very dark skin.Please answer 'yes' or 'no' only.",
'images' => [$base64],
};
my $url = "$ollama/api/generate";
my $res = GetRest($url, $ref);
print "$file\t$res\n";
#Remove thought information from the response
$res =~ s/<\|channel>.*?<channel\|>//s;
$res =~ s/^\s+|\s+$//g;
$total++;
if($res=~/yes/i) {
$yes++;
print "\tYES\n";
}
unlink($file);
}
if($total == 0) {
warn "No images were processed for $path\n";
return undef;
}
my $percentage = ($yes*100)/$total;
print "$yes of $total images gave yes - $percentage\n";
my $verdict = 0;
if($percentage >= $sensitivity) {
print "$path: Probably a black dude, $percentage\n";
$verdict = 1;
} else {
print "$path: Probably not, $percentage\n";
$verdict = 0;
}
return {
'images' => $total,
'yes' => $yes,
'percentage'=> $percentage,
'verdict' => $verdict,
'model' => $model,
'path' => $path,
};
}
sub GetRest {
my($url, $ref)=@_;
my $payload = encode_json($ref);
my $response = $ua->post($url, 'Content_Type' => 'application/json', 'Content' => $payload);
if($response->is_success) {
my $content = $response->decoded_content;
my @parts = split(/\n/, $content);
my @refs = map { decode_json($_) } @parts;
return join(' ', map { $_->{'response'} } @refs);
} else {
die "Error: " . $response->status_line;
}
}
sub PostResponse {
my ($task_id, $assign_key, $results) = @_;
my $url = "$api/api/v1/task/$task_id/complete";
my $ref = {
'assign_key' => $assign_key,
'results' => $results,
};
my $payload = encode_json($ref);
my $response = $ua->post($url, 'Content_Type' => 'application/json', 'Content' => $payload);
if($response->is_success) {
return $response->decoded_content;
} else {
die "Error posting response ($url): " . $response->status_line;
}
}
sub GetNext {
my $url = "$api/api/v1/nexttask/AISCAN";
my $task;
my $video;
my $response = $ua->get($url);
if($response->is_success) {
$task=decode_json($response->decoded_content);
} else {
die "Error fetching next task ($url): " . $response->status_line;
}
my $video_id = $task->{'task'}{'video_id'};
$url = "$api/api/v1/video/$video_id";
$response = $ua->get($url);
if($response->is_success) {
$video=decode_json($response->decoded_content);
$task->{'video'} = $video;
} else {
die "Error fetching video ($url): " . $response->status_line;
}
return $task;
}
+84 -21
View File
@@ -3,6 +3,7 @@ use strict;
use Dancer2; use Dancer2;
use Dancer2::Plugin::Database; use Dancer2::Plugin::Database;
use Data::Dumper; use Data::Dumper;
use JSON ();
use SQL::SplitStatement; use SQL::SplitStatement;
use HTTP::Lite; use HTTP::Lite;
@@ -83,12 +84,42 @@ post '/api/v1/video' => sub {
get '/api/v1/tasks' => sub { get '/api/v1/tasks' => sub {
my $status = query_parameters->get("status") || 'PENDING'; my $status = query_parameters->get("status") || 'PENDING';
my $sth = database->prepare("SELECT id, video_id, status FROM tasks WHERE status=? ORDER BY id"); my $type = query_parameters->get("type");
$sth->execute($status);
my $ref = $sth->fetchall_hashref('id'); unless ($type) {
send_error("Missing 'type' query parameter", 400);
}
# Optional pagination
my $page = query_parameters->get("page") || 1;
my $per_page = query_parameters->get("per_page") || 20;
my $offset = ($page - 1) * $per_page;
my $sth = database->prepare(
"SELECT id, video_id, task_type AS type, status FROM tasks WHERE status=? AND task_type=? ORDER BY id LIMIT ? OFFSET ?"
);
$sth->execute($status, $type, $per_page, $offset);
my $ref = [];
while(my $row = $sth->fetchrow_hashref()) {
push @$ref, $row;
}
$sth->finish(); $sth->finish();
return $ref; # Total count for pagination metadata
my $count_sth = database->prepare(
"SELECT COUNT(*) AS cnt FROM tasks WHERE status=? AND task_type=?"
);
$count_sth->execute($status, $type);
my $cnt = ($count_sth->fetchrow_hashref() || { cnt => 0 })->{cnt};
$count_sth->finish();
return {
tasks => $ref,
page => $page,
per_page => $per_page,
total => $cnt,
total_pages => max(1, int(($cnt + $per_page - 1) / $per_page)),
};
}; };
get '/api/v1/task/:task' => sub { get '/api/v1/task/:task' => sub {
@@ -107,7 +138,7 @@ get '/api/v1/task/:task' => sub {
get '/api/v1/nexttask/:type' => sub { get '/api/v1/nexttask/:type' => sub {
my $type = route_parameters->get("type"); my $type = route_parameters->get("type");
my $sth = database->prepare("SELECT * FROM tasks WHERE task_type=? AND status='PENDING' ORDER BY id LIMIT 1"); my $sth = database->prepare("SELECT * FROM tasks WHERE task_type=? AND status='PENDING' ORDER BY updated_at ASC LIMIT 1");
$sth->execute($type); $sth->execute($type);
my $ref = $sth->fetchrow_hashref(); my $ref = $sth->fetchrow_hashref();
$sth->finish(); $sth->finish();
@@ -129,13 +160,13 @@ get '/api/v1/nexttask/:type' => sub {
}; };
post '/api/v1/task/:task/complete' => sub { post '/api/v1/task/:task/complete' => sub {
my $task = route_parameters->get("task"); my $task_id = route_parameters->get("task");
# Accept both JSON body and form-encoded data for assign_key/results. # Accept both JSON body and form-encoded data for assign_key/results.
my ($assign_key, $results); my ($assign_key, $results);
my $ct = request_header('Content-Type') || ''; my $ct = request_header('Content-Type') || '';
if ($ct eq 'application/json') { if ($ct =~ m{application/json}i) {
my $body = decode_json(request->body()); my $body = JSON::decode_json(request->body());
$assign_key = $body->{assign_key}; $assign_key = $body->{assign_key};
$results = $body->{results}; $results = $body->{results};
} else { } else {
@@ -143,24 +174,31 @@ post '/api/v1/task/:task/complete' => sub {
$results = body_parameters->get("results"); $results = body_parameters->get("results");
} }
unless ($assign_key) { # Optional worker verification: if assign_key is provided, require IN_PROGRESS + matching key.
send_error("Missing assign_key", 400); # When omitted (e.g., UI-driven workflows), skip worker checks and just complete the task.
} if ($assign_key) {
my $sth = database->prepare("SELECT * FROM tasks WHERE id=? AND assign_key=? AND status='IN_PROGRESS'");
$sth->execute($task_id, $assign_key);
my $ref = $sth->fetchrow_hashref();
$sth->finish();
# Verify the task is assigned to the worker unless(ref $ref) {
my $sth = database->prepare("SELECT * FROM tasks WHERE id=? AND assign_key=? AND status='IN_PROGRESS'"); send_error("Task not assigned to this worker or not in progress", 403);
$sth->execute($task, $assign_key); }
my $ref = $sth->fetchrow_hashref(); } else {
$sth->finish(); # No assign_key provided — caller may be an unauthenticated UI, etc.
my $sth = database->prepare("SELECT id FROM tasks WHERE id=?");
unless(ref $ref) { $sth->execute($task_id);
send_error("Task not assigned to this worker or not in progress", 403); unless ($sth->fetchrow_hashref()) {
send_error("Task not found", 404);
}
$sth->finish();
} }
# Update the task as completed (store results as JSON text for MariaDB) # Update the task as completed (store results as JSON text for MariaDB)
my $results_json = defined $results ? encode_json($results) : 'null'; my $results_json = defined $results ? JSON::encode_json($results) : 'null';
my $update_sth = database->prepare("UPDATE tasks SET status='COMPLETED', results=?, updated_at=NOW() WHERE id=?"); my $update_sth = database->prepare("UPDATE tasks SET status='COMPLETED', results=?, updated_at=NOW() WHERE id=?");
$update_sth->execute($results_json, $task); $update_sth->execute($results_json, $task_id);
if($update_sth->rows == 0) { if($update_sth->rows == 0) {
send_error("Failed to complete task", 500); send_error("Failed to complete task", 500);
} }
@@ -183,4 +221,29 @@ post '/api/v1/task' => sub {
return { id => $task_id }; return { id => $task_id };
}; };
post '/api/v1/task/:task/unassign' => sub {
my $task_id = route_parameters->get("task");
my $update_sth = database->prepare(
"UPDATE tasks SET status='PENDING', assign_key=NULL, assigned_at=NULL, updated_at=NOW() WHERE id=?"
);
$update_sth->execute($task_id);
if ($update_sth->rows == 0) {
send_error("Task not found or already in PENDING state", 404);
}
return { message => 'Task unassigned' };
};
start(); start();
sub max {
return (sort { $b <=> $a } @_)[0];
}
sub min {
return (sort { $a <=> $b } @_)[0];
}
+2 -2
View File
@@ -13,12 +13,12 @@ plugins:
host: 'mariadb' host: 'mariadb'
port: 3306 port: 3306
username: 'videodetect' username: 'videodetect'
password: 'changeme_videodetect' password: 'videodetect'
connection_check_threshold: 10 connection_check_threshold: 10
dbi_params: dbi_params:
RaiseError: 1 RaiseError: 1
AutoCommit: 1 AutoCommit: 1
on_connect_do: ["SET NAMES 'utf8'", "SET CHARACTER SET 'utf8'" ] on_connect_do: ["SET NAMES 'utf8mb4'", "SET CHARACTER SET 'utf8mb4'" ]
log_queries: 1 log_queries: 1
server: server:
Executable
+2
View File
@@ -0,0 +1,2 @@
#!/bin/bash
docker run --rm -it mariadb:10.11 mariadb --user=videodetect --password=videodetect --host=10.0.1.7 --port=3306 videodetect
+9 -17
View File
@@ -7,7 +7,7 @@ services:
MYSQL_ROOT_PASSWORD: ${DB_ROOT_PASSWORD:-rootpass} MYSQL_ROOT_PASSWORD: ${DB_ROOT_PASSWORD:-rootpass}
MYSQL_DATABASE: ${DB_NAME:-videodetect} MYSQL_DATABASE: ${DB_NAME:-videodetect}
MYSQL_USER: ${DB_USER:-videodetect} MYSQL_USER: ${DB_USER:-videodetect}
MYSQL_PASSWORD: ${DB_PASSWORD:-changeme_videodetect} MYSQL_PASSWORD: ${DB_PASSWORD:-videodetect}
ports: ports:
- "3306:3306" - "3306:3306"
volumes: volumes:
@@ -29,21 +29,14 @@ services:
build: build:
context: ./ui context: ./ui
dockerfile: Dockerfile dockerfile: Dockerfile
container_name: videodetect-ui
restart: unless-stopped restart: unless-stopped
ports: ports:
- "5000:5000" - "5000:5000"
environment: environment:
- DB_HOST=mariadb - API_BASE=http://10.0.1.34:3000/api/v1
- DB_PORT=3306
- DB_NAME=videodetect
- DB_USER=videodetect
- DB_PASSWORD=${DB_PASSWORD:-changeme_videodetect}
- FLASK_ENV=production
volumes: volumes:
- ./ui:/app - ./ui:/app
- nas_input:/data/input:ro - nas_input:/data:ro
- ${NAS_OUTPUT_PATH:-/mnt/nas/output}:/data/output
networks: networks:
- videodetect-network - videodetect-network
depends_on: depends_on:
@@ -59,16 +52,15 @@ services:
context: ./api context: ./api
dockerfile: Dockerfile dockerfile: Dockerfile
container_name: videodetect-api container_name: videodetect-api
restart: no restart: unless-stopped
ports: ports:
- "3000:3000" - "3000:3000"
environment: environment:
- DB_HOST=mariadb - DB_HOST=${DB_HOST:-mariadb}
- DB_PORT=3306 - DB_PORT=${DB_PORT:-3306}
- DB_NAME=videodetect - DB_NAME=${DB_NAME:-videodetect}
- DB_USER=videodetect - DB_USER=${DB_USER:-videodetect}
- DB_PASSWORD=${DB_PASSWORD:-changeme_videodetect} - DB_PASSWORD=${DB_PASSWORD:-videodetect}
- DANCER_ENVIRONMENT=production
volumes: volumes:
- ./api:/app - ./api:/app
networks: networks:
Executable
+5
View File
@@ -0,0 +1,5 @@
#!/bin/bash
curl http://10.0.1.7:11434/api/generate -d '{
"model": "qwen3.6:35b-a3b",
"keep_alive": "60m"
}'
-68
View File
@@ -1,68 +0,0 @@
#!/usr/bin/perl
use strict;
use MIME::Base64 qw(encode_base64);
use LWP::UserAgent;
use JSON;
use Data::Dumper;
my $seconds = 30;
my $ollama = 'http://10.0.1.7:11434';
my $resolution = "672:672";
my $sensitivity = 20;
my $file = $ARGV[0] || die "Usage: $0 <filename>\n";
system("ffmpeg -i \"$file\" -vf fps=1/$seconds,scale=$resolution output_%04d.png");
my $ua = LWP::UserAgent->new;
$ua->timeout(60);
my $total=0;
my $yes=0;
$/ = undef;
foreach my $file (<*.png>) {
print "$file...\n";
open(IN, '<:raw', $file) || die "Unable to read $file: $!\n";
my $data = <IN>;
close(IN);
my $base64 = encode_base64($data);
$data=undef;
my $ref = {
'model' => 'srizon/pixie:latest',
'prompt' => "Are there any african american people or in this photo, whether they are shown completely or not? Please answer 'yes' or 'no' only.",
'images' => [$base64],
};
my $url = "$ollama/api/generate";
my $res = GetRest($url, $ref);
print "\t$res\n";
$total++;
if($res=~/yes/i) {
$yes++;
}
unlink($file);
}
my $percentage = ($yes*100)/$total;
print "$yes of $total images gave yes - $percentage\n";
if($percentage >= $sensitivity) {
print "Probably a black dude, $percentage\n";
} else {
print "Probably not, $percentage\n";
}
sub GetRest {
my($url, $ref)=@_;
my $payload = encode_json($ref);
my $response = $ua->post($url, 'Content_Type' => 'application/json', 'Content' => $payload);
if($response->is_success) {
my $content = $response->decoded_content;
my @parts = split(/\n/, $content);
my @refs = map { decode_json($_) } @parts;
return join(' ', map { $_->{'response'} } @refs);
} else {
die "Error: " . $response->status_line;
}
}
+14
View File
@@ -0,0 +1,14 @@
# VideoDetect — Docker Compose environment variables
# Copy this file to .env and fill in real values before running docker-compose up.
# -----------------------------------------------------
# Database
# -----------------------------------------------------
DB_ROOT_PASSWORD=videodetect
DB_PASSWORD=videodetect
DB_USER=videodetect
DB_NAME=videodetect
DB_HOST=10.0.1.34
API_HOST=https://10.0.1.34:3000
-6
View File
@@ -18,8 +18,6 @@ services:
volumes: volumes:
- nas_input:/data:ro - nas_input:/data:ro
- $PWD/scanner.pl:/app/scanner.pl:ro - $PWD/scanner.pl:/app/scanner.pl:ro
networks:
- videodetect_videodetect-network
volumes: volumes:
nas_input: nas_input:
@@ -28,7 +26,3 @@ volumes:
type: nfs type: nfs
o: addr=10.0.0.2,ro,nfsvers=4,hard,intr o: addr=10.0.0.2,ro,nfsvers=4,hard,intr
device: ":/mnt/Bulk/Homes/ryan/Prawns" device: ":/mnt/Bulk/Homes/ryan/Prawns"
networks:
videodetect_videodetect-network:
external: true
+74 -22
View File
@@ -9,14 +9,23 @@ use String::ShellQuote qw/shell_quote/;
my $db_host = $ENV{'DB_HOST'} || 'mariadb'; my $db_host = $ENV{'DB_HOST'} || 'mariadb';
my $db_name = $ENV{'DB_NAME'} || 'videodetect'; my $db_name = $ENV{'DB_NAME'} || 'videodetect';
my $db_user = $ENV{'DB_USER'} || 'videodetect'; my $db_user = $ENV{'DB_USER'} || 'videodetect';
my $db_pass = $ENV{'DB_PASSWORD'} || 'videodetect123'; my $db_pass = $ENV{'DB_PASSWORD'} || 'videodetect';
my $do_scan = (defined $ENV{'DO_SCAN'}) ? $ENV{'DO_SCAN'} : 1;
my $do_review = (defined $ENV{'DO_REVIEW'}) ? $ENV{'DO_REVIEW'} : 1;
my $do_delete = (defined $ENV{'DO_DELETE'}) ? $ENV{'DO_DELETE'} : 1;
$SIG{INT} = sub { die "Interrupted\n"; }; $SIG{INT} = sub { die "Interrupted\n"; };
my $dbh = DBI->connect("DBI:mysql:database=$db_name;host=$db_host", $db_user, $db_pass); my $dbh = DBI->connect(
"DBI:mysql:database=$db_name;host=$db_host;mysql_enable_utf8mb4=1",
$db_user,
$db_pass,
{ mysql_auto_reconnect => 1 }
);
my @video_extensions = qw(mp4 mkv avi mov flv wmv mpg mpeg webm); my @video_extensions = qw(mp4 mkv avi mov flv wmv mpg mpeg webm m4v);
my %ignored_extensions;
my @dir_queue = ('/data'); my @dir_queue = ('/data');
@@ -27,35 +36,46 @@ my $row = $sth->fetchrow_hashref();
my $now = $row->{now}; my $now = $row->{now};
$sth->finish(); $sth->finish();
while(my $dir = shift @dir_queue) { if($do_scan) {
opendir(my $dh, $dir) or die "Cannot open directory $dir: $!"; while(my $dir = shift @dir_queue) {
print "$dir\n"; opendir(my $dh, $dir) or die "Cannot open directory $dir: $!";
while (my $file = readdir($dh)) { print "$dir\n";
next if ($file eq '.' || $file eq '..'); while (my $file = readdir($dh)) {
my $full_path = "$dir/$file"; next if ($file eq '.' || $file eq '..');
if (-d $full_path) { my $full_path = "$dir/$file";
unshift @dir_queue, $full_path; if (-d $full_path) {
} elsif (-f $full_path) { unshift @dir_queue, $full_path;
process_file($full_path); } elsif (-f $full_path) {
process_file($full_path);
}
} }
closedir($dh);
}
for my $ext (sort { $ignored_extensions{$b} <=> $ignored_extensions{$a} } keys %ignored_extensions) {
print "Ignored extension $ext: $ignored_extensions{$ext} times\n";
} }
closedir($dh);
} }
# Any entry where last_scan_time < $now is considered deleted or moved, so we can mark them as such if($do_scan && $do_delete) {
my $sth_delete = $dbh->prepare("DELETE FROM videos WHERE last_scan_time < ? OR last_scan_time IS NULL"); # Any entry where last_scan_time < $now is considered deleted or moved, so we can mark them as such
$sth_delete->execute($now); my $sth_delete = $dbh->prepare("DELETE FROM videos WHERE last_scan_time < ? OR last_scan_time IS NULL");
$sth_delete->execute($now);
}
if($do_review) {
print "Scanning complete. Creating review tasks for completed AISCAN tasks.\n";
CreateReviewTasks();
}
sub process_file { sub process_file {
my ($file_path) = @_; my ($file_path) = @_;
# Determine if the file is a video based on mime info or extension. For simplicity, let's check the extension.
if (my ($ext)=$file_path =~ /^.+\.(\S+?)$/) { if (my ($ext)=$file_path =~ /^.+\.(\S+?)$/) {
unless (grep { lc($ext) eq $_ } @video_extensions) { unless (grep { lc($ext) eq $_ } @video_extensions) {
$ignored_extensions{lc($ext)}++;
return; # Not a video file return; # Not a video file
} }
} }
print "$file_path\n";
my $sth = $dbh->prepare("SELECT id,file_size FROM videos WHERE file_path=?"); my $sth = $dbh->prepare("SELECT id,file_size FROM videos WHERE file_path=?");
$sth->execute($file_path); $sth->execute($file_path);
if(my $row = $sth->fetchrow_hashref()) { if(my $row = $sth->fetchrow_hashref()) {
@@ -66,9 +86,15 @@ sub process_file {
$update_sth->execute($file_size, $row->{id}); $update_sth->execute($file_size, $row->{id});
$update_sth->finish(); $update_sth->finish();
create_aiscan_task($row->{id}); create_aiscan_task($row->{id});
} else {
my $update_sth = $dbh->prepare("UPDATE videos SET last_scan_time=NOW() WHERE id=?");
$update_sth->execute($row->{id});
$update_sth->finish();
} }
return; # Already exists and size matches return; # Already exists and size matches
} }
print "NEW:$file_path\n";
my $info = get_video_info($file_path); my $info = get_video_info($file_path);
unless ($info) { unless ($info) {
@@ -93,9 +119,14 @@ sub process_file {
sub create_aiscan_task { sub create_aiscan_task {
my ($video_id) = @_; my ($video_id) = @_;
$dbh->do("DELETE FROM tasks WHERE video_id=$video_id AND task_type='AISCAN'"); create_task($video_id, 'AISCAN');
my $sth = $dbh->prepare("INSERT INTO tasks (video_id, task_type, status) VALUES (?, 'AISCAN', 'PENDING')"); }
$sth->execute($video_id);
sub create_task {
my ($video_id, $task_type) = @_;
$dbh->do("DELETE FROM tasks WHERE video_id=$video_id AND task_type=?", undef, $task_type);
my $sth = $dbh->prepare("INSERT INTO tasks (video_id, task_type, status) VALUES (?, ?, 'PENDING')");
$sth->execute($video_id, $task_type);
$sth->finish(); $sth->finish();
} }
@@ -159,3 +190,24 @@ sub get_video_info {
duration => $duration, duration => $duration,
}; };
} }
sub CreateReviewTasks {
# Fetch all completed AISCAN tasks and filter in Perl.
# This avoids MariaDB's inconsistent JSON functions in SQL WHERE clauses,
# which behave differently between CLI client and DBI connections.
my $sth = $dbh->prepare("
SELECT id, video_id
FROM tasks
WHERE task_type='AISCAN'
AND status='COMPLETED'
AND JSON_EXTRACT(results, '\$.verdict') = 1
AND (SELECT 1 FROM tasks t2 WHERE t2.video_id = tasks.video_id AND t2.task_type = 'REVIEW' LIMIT 1) IS NULL
");
$sth->execute();
while (my $row = $sth->fetchrow_hashref()) {
print "Creating REVIEW task for video_id: $row->{video_id}\n";
create_task($row->{video_id}, 'REVIEW');
}
$sth->finish();
}
-1
View File
@@ -1 +0,0 @@
"""Active learning pipeline: label ingestion, training, validation, and deployment."""
-217
View File
@@ -1,217 +0,0 @@
"""
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 []
-82
View File
@@ -1,82 +0,0 @@
"""
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
-223
View File
@@ -1,223 +0,0 @@
"""
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)
-344
View File
@@ -1,344 +0,0 @@
"""
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
-204
View File
@@ -1,204 +0,0 @@
"""
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
-39
View File
@@ -1,39 +0,0 @@
"""Confidence aggregation strategies for video-level scoring."""
import logging
from typing import List
import numpy as np
logger = logging.getLogger(__name__)
def aggregate(
confidences: List[float],
strategy: str = "max",
alpha: float = 1.0,
beta: float = 0.1,
top_k: int = 3,
) -> float:
"""Aggregate per-crop confidences into a single video-level score C ∈ [0,1]."""
if not confidences:
return 0.0
arr = np.array(confidences, dtype=np.float64)
if strategy == "max":
return float(arr.max())
if strategy == "weighted_mean":
mean = float(arr.mean())
var = float(arr.var())
raw = alpha * mean + beta * var
# Clamp to [0,1]; the formula is a linear combination, not a softmax
return float(np.clip(raw, 0.0, 1.0))
if strategy == "top_k_mean":
k = max(1, min(top_k, len(arr)))
return float(np.partition(arr, -k)[-k:].mean())
logger.warning("Unknown aggregation strategy '%s', falling back to max", strategy)
return float(arr.max())
-136
View File
@@ -1,136 +0,0 @@
"""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()
-68
View File
@@ -1,68 +0,0 @@
"""Dynamic batching utilities for face detection inference."""
import logging
import time
from collections import deque
from typing import Callable, Deque, Generic, List, Optional, TypeVar
from gpu_manager import GPUMemoryManager
logger = logging.getLogger(__name__)
T = TypeVar("T")
R = TypeVar("R")
class DynamicBatcher(Generic[T, R]):
"""Accumulate items into batches and flush based on size or timeout."""
def __init__(
self,
process_fn: Callable[[List[T]], List[R]],
gpu_manager: Optional[GPUMemoryManager] = None,
max_batch_size: int = 16,
batch_timeout_ms: float = 100.0,
min_batch_size: int = 1,
):
self.process_fn = process_fn
self.gpu_manager = gpu_manager
self.max_batch_size = max_batch_size
self.batch_timeout_ms = batch_timeout_ms
self.min_batch_size = min_batch_size
self._queue: Deque[T] = deque()
self._last_flush = time.monotonic()
def add(self, item: T) -> List[R]:
"""Add an item and return any results if a batch was flushed."""
self._queue.append(item)
if len(self._queue) >= self._current_batch_size():
return self.flush()
elapsed_ms = (time.monotonic() - self._last_flush) * 1000
if elapsed_ms >= self.batch_timeout_ms and len(self._queue) >= self.min_batch_size:
return self.flush()
return []
def flush(self) -> List[R]:
"""Flush all queued items through the processor."""
if not self._queue:
return []
batch = [self._queue.popleft() for _ in range(min(len(self._queue), self._current_batch_size()))]
results = self.process_fn(batch)
self._last_flush = time.monotonic()
if self.gpu_manager is not None:
self.gpu_manager.adjust_batch_size()
self.gpu_manager.empty_cache()
return results
def _current_batch_size(self) -> int:
if self.gpu_manager is not None:
return min(self.gpu_manager.current_batch_size, self.max_batch_size)
return self.max_batch_size
@property
def queued_count(self) -> int:
return len(self._queue)
-138
View File
@@ -1,138 +0,0 @@
"""
Face crop classifier.
Loads a MobileNetV3 model via ONNX/TensorRT FP32, applies temperature-scaled
softmax, and returns per-crop probabilities for the target class.
"""
import logging
from pathlib import Path
from typing import Any, List, Optional
import numpy as np
logger = logging.getLogger(__name__)
# ImageNet normalisation constants
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
def calibrated_softmax(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
"""Softmax with temperature scaling; returns class probabilities."""
scaled = logits / max(temperature, 1e-8)
shifted = scaled - scaled.max(axis=-1, keepdims=True) # numerical stability
exp = np.exp(shifted)
return exp / exp.sum(axis=-1, keepdims=True)
class FaceClassifier:
"""Classify face crops and return calibrated target-class probabilities."""
def __init__(
self,
engine_path: str,
temperature: float = 1.0,
input_size: int = 224,
device: str = "cuda",
):
self.engine_path = engine_path
self.temperature = temperature
self.input_size = input_size
self.device = device
self._session: Optional[Any] = None
self._load_model()
def _load_model(self):
path = Path(self.engine_path)
if not path.exists():
logger.warning("Classifier engine not found at %s; using placeholder", self.engine_path)
return
suffix = path.suffix.lower()
if suffix == ".onnx":
try:
import onnxruntime as ort
providers = (["CUDAExecutionProvider"] if self.device.startswith("cuda")
else ["CPUExecutionProvider"])
self._session = ort.InferenceSession(str(path), providers=providers)
logger.info("Loaded ONNX classifier from %s", self.engine_path)
except Exception as exc:
logger.warning("Failed to load ONNX classifier: %s", exc)
elif suffix in (".trt", ".engine", ".plan"):
try:
import tensorrt as trt
with trt.Logger() as trt_logger, open(path, "rb") as f:
runtime = trt.Runtime(trt_logger)
self._session = runtime.deserialize_cuda_engine(f.read())
logger.info("Loaded TensorRT classifier from %s", self.engine_path)
except Exception as exc:
logger.warning("Failed to load TensorRT classifier: %s", exc)
else:
logger.warning("Unsupported classifier format: %s", suffix)
def classify(self, crop_paths: List[str], batch_size: int = 16) -> List[float]:
"""Return target-class probabilities (p ∈ [0,1]) for each crop path."""
results: List[float] = []
for i in range(0, len(crop_paths), batch_size):
batch = crop_paths[i : i + batch_size]
results.extend(self._classify_batch(batch))
return results
def _classify_batch(self, crop_paths: List[str]) -> List[float]:
preprocessed = []
for path in crop_paths:
try:
preprocessed.append(self._preprocess(path))
except Exception as exc:
logger.warning("Preprocessing failed for %s: %s", path, exc)
preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
batch = np.stack(preprocessed, axis=0) # (N, 3, H, W)
logits = self._infer(batch) # (N, 2)
probs = calibrated_softmax(logits, self.temperature)
return probs[:, 1].tolist() # target-class column
def _preprocess(self, crop_path: str) -> np.ndarray:
from PIL import Image
image = Image.open(crop_path).convert("RGB")
if image.size != (self.input_size, self.input_size):
image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
arr = np.array(image, dtype=np.float32) / 255.0
arr = (arr - _MEAN) / _STD
image.close()
return np.transpose(arr, (2, 0, 1)) # HWC → CHW
def _infer(self, batch: np.ndarray) -> np.ndarray:
if self._session is None:
return self._placeholder_logits(batch.shape[0])
try:
if hasattr(self._session, "run"):
input_name = self._session.get_inputs()[0].name
output = self._session.run(None, {input_name: batch})
return np.array(output[0])
# TensorRT path
import pycuda.autoinit # noqa: F401
import pycuda.driver as cuda
context = self._session.create_execution_context()
out_shape = (batch.shape[0], 2)
d_in = cuda.mem_alloc(batch.nbytes)
d_out = cuda.mem_alloc(np.prod(out_shape) * np.dtype(np.float32).itemsize)
stream = cuda.Stream()
cuda.memcpy_htod_async(d_in, batch, stream)
context.execute_async_v2(bindings=[int(d_in), int(d_out)], stream_handle=stream.handle)
output = np.empty(out_shape, dtype=np.float32)
cuda.memcpy_dtoh_async(output, d_out, stream)
stream.synchronize()
return output
except Exception as exc:
logger.warning("Classifier inference failed: %s", exc)
return self._placeholder_logits(batch.shape[0])
@staticmethod
def _placeholder_logits(n: int) -> np.ndarray:
"""Return neutral logits when no model is loaded."""
return np.zeros((n, 2), dtype=np.float32)
-107
View File
@@ -1,107 +0,0 @@
"""
Codec Validator
Validates video codecs against a whitelist/blacklist.
Flags unsupported codecs as UNSCANNABLE.
"""
import logging
from typing import Set
logger = logging.getLogger(__name__)
# Default supported codecs
DEFAULT_WHITELIST: Set[str] = {
"avc1", # H.264
"hevc", # H.265
"vp8", # VP8
"vp9", # VP9
"av01", # AV1
"mjpeg", # MJPEG
"mp4v", # MPEG-4
"h264", # H.264 (alternative name)
"h265", # H.265 (alternative name)
}
# Known unsupported codecs
DEFAULT_BLACKLIST: Set[str] = {
"theora", # Theora
"divx", # DivX
"xvid", # Xvid
"prores", # ProRes (requires special handling)
"avc", # Old AVC naming
"hev1", # HEVC alternative
}
class CodecValidator:
"""Validate video codecs against whitelist/blacklist."""
def __init__(
self,
whitelist: Set[str] = None,
blacklist: Set[str] = None,
default_status_on_error: str = "UNSCANNABLE",
):
self.whitelist = whitelist or DEFAULT_WHITELIST
self.blacklist = blacklist or DEFAULT_BLACKLIST
self.default_status_on_error = default_status_on_error
def is_supported(self, codec: str) -> bool:
"""Check if a codec is supported."""
codec_lower = codec.lower().strip()
# Check blacklist first
if codec_lower in self.blacklist:
logger.info("Codec '%s' is blacklisted", codec)
return False
# Check whitelist
if codec_lower in self.whitelist:
return True
# Not in either list - assume unsupported
logger.warning("Codec '%s' not in whitelist or blacklist", codec)
return False
def validate(self, codec: str) -> tuple:
"""
Validate a codec and return (is_supported, reason).
Returns:
(True, "supported") if codec is in whitelist
(False, "blacklisted") if codec is in blacklist
(False, "not_in_whitelist") if codec is unknown
"""
codec_lower = codec.lower().strip()
if codec_lower in self.blacklist:
return False, "blacklisted"
if codec_lower in self.whitelist:
return True, "supported"
return False, "not_in_whitelist"
def add_to_whitelist(self, codec: str):
"""Add a codec to the whitelist."""
self.whitelist.add(codec.lower().strip())
logger.info("Added codec '%s' to whitelist", codec)
def add_to_blacklist(self, codec: str):
"""Add a codec to the blacklist."""
self.blacklist.add(codec.lower().strip())
logger.info("Added codec '%s' to blacklist", codec)
def get_status(self, codec: str) -> str:
"""Get the status for a codec (for DB update)."""
is_supported, reason = self.validate(codec)
if is_supported:
return "PENDING"
return self.default_status_on_error
def __repr__(self):
return (
f"CodecValidator(whitelist={len(self.whitelist)}, "
f"blacklist={len(self.blacklist)})"
)
-126
View File
@@ -1,126 +0,0 @@
"""
Configuration loader for VideoDetect.
Loads config.yaml with environment variable overrides.
Provides typed accessors and validates all required fields.
"""
import logging
import os
from pathlib import Path
from typing import Any, Dict, Optional
import yaml
logger = logging.getLogger(__name__)
# Default config path
DEFAULT_CONFIG_PATH = Path(__file__).parent.parent / "config.yaml"
class Config:
"""Typed configuration manager with environment variable overrides."""
def __init__(self, config_path: Optional[str] = None):
self._path = Path(config_path) if config_path else DEFAULT_CONFIG_PATH
self._data: Dict[str, Any] = {}
self._load()
def _load(self):
"""Load config from YAML file."""
if not self._path.exists():
raise FileNotFoundError(f"Config file not found: {self._path}")
with open(self._path, "r") as f:
self._data = yaml.safe_load(f) or {}
# Apply environment variable overrides
self._apply_env_overrides()
logger.info("Config loaded from %s", self._path)
def _apply_env_overrides(self):
"""Override config values with environment variables."""
overrides = {
"DB_HOST": ("database.host", None),
"DB_PORT": ("database.port", int),
"DB_NAME": ("database.name", None),
"DB_USER": ("database.user", None),
"DB_PASSWORD": ("database.password", None),
"DB_POOL_SIZE": ("database.pool_size", int),
"DB_POOL_MIN": ("database.pool_min", int),
"GPU_MAX_MEMORY_GB": ("gpu.max_memory_gb", float),
"SAMPLING_INTERVAL_SECONDS": ("sampling.interval_seconds", int),
"T_HIGH": ("thresholds.T_high", float),
"T_LOW": ("thresholds.T_low", float),
"CONFIG_PATH": (None, None), # handled separately
}
for env_key, (config_path, type_fn) in overrides.items():
env_val = os.environ.get(env_key)
if env_val is not None:
if config_path is None:
continue
if type_fn is not None:
env_val = type_fn(env_val)
self._set_nested(self._data, config_path, env_val)
logger.debug("Config override: %s=%s (from %s)", config_path, env_val, env_key)
@staticmethod
def _set_nested(data: Dict, path: str, value):
"""Set a value in a nested dict using dot notation."""
keys = path.split(".")
d = data
for key in keys[:-1]:
d = d.setdefault(key, {})
d[keys[-1]] = value
def get(self, path: str, default=None):
"""Get a config value using dot notation."""
keys = path.split(".")
d = self._data
for key in keys:
if isinstance(d, dict):
d = d.get(key, default)
else:
return default
return d if d is not None else default
def get_section(self, section: str) -> Dict:
"""Get an entire config section as a dict."""
return self._data.get(section, {})
def validate(self, required_keys: list) -> list:
"""Validate that required config keys exist. Returns list of missing keys."""
missing = []
for key in required_keys:
if self.get(key) is None:
missing.append(key)
if missing:
logger.error("Missing required config keys: %s", missing)
return missing
@property
def data(self) -> Dict[str, Any]:
"""Access raw config data."""
return self._data
def __repr__(self):
return f"Config(path={self._path}, sections={list(self._data.keys())})"
# Module-level singleton
_config: Optional[Config] = None
def get_config(config_path: Optional[str] = None) -> Config:
"""Get or create the global config singleton."""
global _config
if _config is None:
_config = Config(config_path)
return _config
def reset_config():
"""Reset the config singleton (useful for testing)."""
global _config
_config = None
-137
View File
@@ -1,137 +0,0 @@
"""
Crash recovery and idempotent processing guarantees.
On startup: re-queues any video stuck in PROCESSING past the lock timeout.
During processing: saves per-video checkpoint JSON to scratch for resume.
"""
import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional
logger = logging.getLogger(__name__)
_CHECKPOINT_DIR = "checkpoints"
class CrashRecovery:
"""Handle crash recovery, checkpointing, and idempotent re-queueing."""
def __init__(self, db_connector, config):
self._db = db_connector
scratch = config.get("storage.scratch_path", "/scratch")
self._checkpoint_dir = Path(scratch) / _CHECKPOINT_DIR
mon = config.get_section("monitoring")
cr = mon.get("crash_recovery", {})
self._lock_timeout_minutes: int = int(cr.get("lock_timeout_minutes", 5))
self._auto_requeue: bool = bool(cr.get("auto_requeue", True))
# ------------------------------------------------------------------
# Startup recovery
# ------------------------------------------------------------------
def recover_on_startup(self) -> int:
"""
Re-queue videos stuck in PROCESSING past the lock timeout.
Returns the count of videos re-queued.
"""
if not self._auto_requeue:
logger.info("auto_requeue disabled; skipping crash recovery")
return 0
try:
result = self._db.execute(
"""UPDATE videos
SET status = 'PENDING', updated_at = NOW()
WHERE status = 'PROCESSING'
AND updated_at < NOW() - INTERVAL %s MINUTE""",
(self._lock_timeout_minutes,),
)
count = result if isinstance(result, int) else 0
if count:
logger.warning("Crash recovery: re-queued %d stuck PROCESSING videos", count)
else:
logger.info("Crash recovery: no stuck videos found")
return count
except Exception as exc:
logger.error("Crash recovery query failed: %s", exc)
return 0
def list_stuck_videos(self) -> List[Dict]:
"""Return videos currently stuck past the lock timeout (read-only)."""
try:
rows = self._db.fetchall(
"""SELECT id, file_path, updated_at
FROM videos
WHERE status = 'PROCESSING'
AND updated_at < NOW() - INTERVAL %s MINUTE""",
(self._lock_timeout_minutes,),
)
return rows or []
except Exception as exc:
logger.error("Failed to list stuck videos: %s", exc)
return []
# ------------------------------------------------------------------
# Per-video checkpointing
# ------------------------------------------------------------------
def save_checkpoint(self, video_id: int, state: str, progress: Dict) -> None:
"""Persist a checkpoint JSON for the given video."""
self._checkpoint_dir.mkdir(parents=True, exist_ok=True)
checkpoint = {
"video_id": video_id,
"state": state,
"progress": progress,
"timestamp": datetime.now(timezone.utc).isoformat(),
}
path = self._checkpoint_dir / f"{video_id}.json"
try:
path.write_text(json.dumps(checkpoint, indent=2), encoding="utf-8")
logger.debug("Saved checkpoint for video %d (state=%s)", video_id, state)
except Exception as exc:
logger.warning("Failed to save checkpoint for video %d: %s", video_id, exc)
def load_checkpoint(self, video_id: int) -> Optional[Dict]:
"""Return the stored checkpoint for video_id, or None if absent."""
path = self._checkpoint_dir / f"{video_id}.json"
if not path.exists():
return None
try:
return json.loads(path.read_text(encoding="utf-8"))
except Exception as exc:
logger.warning("Failed to read checkpoint for video %d: %s", video_id, exc)
return None
def delete_checkpoint(self, video_id: int) -> None:
"""Remove the checkpoint file after successful processing."""
path = self._checkpoint_dir / f"{video_id}.json"
if path.exists():
try:
path.unlink()
logger.debug("Deleted checkpoint for video %d", video_id)
except Exception as exc:
logger.warning("Failed to delete checkpoint for video %d: %s", video_id, exc)
# ------------------------------------------------------------------
# Idempotency guard
# ------------------------------------------------------------------
def is_already_completed(self, video_id: int) -> bool:
"""
Return True if this video already has a COMPLETED status.
Prevents duplicate processing if a job is accidentally re-submitted.
"""
try:
rows = self._db.fetchall(
"SELECT status FROM videos WHERE id = %s",
(video_id,),
)
return bool(rows and rows[0].get("status") == "COMPLETED")
except Exception as exc:
logger.warning("Idempotency check failed for video %d: %s", video_id, exc)
return False
-115
View File
@@ -1,115 +0,0 @@
"""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
-227
View File
@@ -1,227 +0,0 @@
"""
Database connection layer with connection pooling for VideoDetect.
Provides:
- Connection pooling via DBUtils + PyMySQL
- Automatic reconnection on disconnect (ping + retry on transient errors)
- Context manager support
- Prepared statements for all queries
- Transaction support for atomic state transitions
"""
import contextlib
import logging
import time
from typing import Optional
from dbutils.pooled_db import PooledDB
import pymysql
logger = logging.getLogger(__name__)
class DBConnector:
"""Thread-safe database connection pool manager."""
# Transient errors that indicate a dropped/stale connection and are safe
# to retry with a fresh connection from the pool.
_RETRYABLE_ERRORS = (pymysql.err.InterfaceError, pymysql.err.OperationalError)
_MAX_RETRIES = 3
_RETRY_BACKOFF_SECONDS = 0.5
def __init__(
self,
host: str = "mariadb",
port: int = 3306,
database: str = "videodetect",
user: str = "videodetect",
password: str = "videodetect123",
pool_size: int = 20,
pool_min: int = 5,
pool_recycle: int = 3600,
):
# PooledDB parameters
pool_config = {
"creator": pymysql,
"maxconnections": pool_size,
"mincached": pool_min,
"maxcached": pool_size,
"maxusage": 200,
"blocking": True,
# ping(0) = check if connection is alive BEFORE returning it from
# the pool. This is the correct DBUtils mechanism for detecting
# dead/stale connections (idle timeout, network blip, DB restart).
"ping": 0,
}
# PyMySQL connection parameters
connection_params = {
"host": host,
"port": port,
"database": database,
"user": user,
"password": password,
"charset": "utf8mb4",
"cursorclass": pymysql.cursors.DictCursor,
"read_timeout": 30,
"write_timeout": 30,
}
self._pool = PooledDB(**pool_config, **connection_params)
self._pool_recycle = pool_recycle
logger.info(
"DB pool initialized: host=%s db=%s pool_size=%d min=%d ping=0",
host, database, pool_size, pool_min,
)
def get_connection(self):
"""Get a connection from the pool."""
return self._pool.connection()
def _run_query(self, query: str, params, transaction: bool, fetch: str):
"""Run a single query on a fresh connection, with retry on transient errors.
A pooled connection can be dropped by the server (idle timeout, network
blip, DB restart). When that happens the query fails with an
InterfaceError/OperationalError. We discard the dead connection and
retry with a fresh one from the pool.
Args:
query: SQL statement.
params: Bound parameters (or None).
transaction: If True, wrap the query in a transaction.
fetch: One of "none", "one", or "all" controlling the return value.
Returns:
- "none": the affected row count.
- "one": a single row (dict) or None.
- "all": a list of rows (dicts).
"""
last_exc = None
for attempt in range(1, self._MAX_RETRIES + 1):
conn = self.get_connection()
try:
cursor = conn.cursor()
if transaction:
conn.begin()
cursor.execute(query, params or ())
if fetch == "one":
result = cursor.fetchone()
elif fetch == "all":
result = cursor.fetchall()
else:
result = cursor.rowcount
if transaction:
conn.commit()
cursor.close()
conn.close()
return result
except self._RETRYABLE_ERRORS as e:
last_exc = e
logger.warning(
"Transient DB error (attempt %d/%d): %s",
attempt, self._MAX_RETRIES, e,
)
# Discard the dead connection; a fresh one is taken on retry.
self._safe_close(conn)
if attempt < self._MAX_RETRIES:
time.sleep(self._RETRY_BACKOFF_SECONDS * attempt)
except Exception:
# Non-transient error: roll back if in a transaction and re-raise.
if transaction:
self._safe_rollback(conn)
self._safe_close(conn)
raise
raise last_exc
@staticmethod
def _safe_close(conn):
try:
conn.close()
except Exception:
pass
@staticmethod
def _safe_rollback(conn):
try:
conn.rollback()
except Exception:
pass
@contextlib.contextmanager
def get_cursor(self, transaction: bool = False):
"""Context manager for getting a cursor with optional transaction support.
Note: this does not retry on transient errors; use execute/fetchone/
fetchall for retry-safe single-statement queries.
"""
conn = self.get_connection()
cursor = None
try:
cursor = conn.cursor()
if transaction:
conn.begin()
yield cursor
if transaction:
conn.commit()
except Exception:
if transaction:
self._safe_rollback(conn)
raise
finally:
if cursor is not None:
try:
cursor.close()
except Exception:
pass
self._safe_close(conn)
@contextlib.contextmanager
def transaction(self):
"""Context manager for a full transaction."""
conn = self.get_connection()
try:
conn.begin()
yield conn
conn.commit()
except Exception:
self._safe_rollback(conn)
raise
finally:
self._safe_close(conn)
def execute(self, query: str, params=None, transaction: bool = False):
"""Execute a query and return affected rows. Retries on transient errors."""
return self._run_query(query, params, transaction, fetch="none")
def fetchone(self, query: str, params=None):
"""Execute a query and return one row. Retries on transient errors."""
return self._run_query(query, params, transaction=False, fetch="one")
def fetchall(self, query: str, params=None):
"""Execute a query and return all rows. Retries on transient errors."""
return self._run_query(query, params, transaction=False, fetch="all")
def initialize_schema(self, schema_path: str = "db/schema.sql"):
"""Initialize the database schema from SQL file."""
with open(schema_path, "r") as f:
sql = f.read()
# Split on semicolons and execute each statement
statements = [s.strip() for s in sql.split(";") if s.strip()]
for stmt in statements:
if stmt.startswith("--"):
continue
try:
self.execute(stmt)
except Exception as e:
logger.debug("Schema statement skipped (may already exist): %s", e)
logger.info("Schema initialized from %s", schema_path)
def health_check(self) -> bool:
"""Check if the database is reachable."""
try:
result = self.fetchone("SELECT 1")
return result is not None
except Exception as e:
logger.error("Health check failed: %s", e)
return False
-219
View File
@@ -1,219 +0,0 @@
"""
Confidence distribution drift detection.
Compares recent video confidence scores against a stored baseline and fires
alerts when the distribution shifts beyond configured thresholds.
"""
import logging
from datetime import datetime, timedelta, timezone
from typing import Callable, Dict, List, Optional
logger = logging.getLogger(__name__)
class DriftDetector:
"""Detect and alert on confidence distribution drift."""
def __init__(self, db_connector, config, alert_fn: Optional[Callable[[str], None]] = None):
self._db = db_connector
mon = config.get_section("monitoring")
alerts = mon.get("alerts", {})
self._drift_threshold: float = float(alerts.get("confidence_drift_threshold", 0.10))
self._review_max_size: int = int(alerts.get("review_queue_max_size", 1000))
self._review_max_age_h: int = int(alerts.get("review_queue_max_age_hours", 24))
self._throughput_min: float = float(alerts.get("throughput_min_videos_per_hour", 20))
self._throughput_window_h: int = int(alerts.get("throughput_min_duration_hours", 1))
self._error_rate_threshold: float = float(alerts.get("error_rate_threshold", 0.05))
self._error_rate_window_h: int = int(alerts.get("error_rate_window_hours", 1))
self._alert_fn = alert_fn or _log_alert
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def run_all_checks(self) -> Dict[str, bool]:
"""
Run all alert checks and return a dict of {check_name: triggered}.
Fires the alert function for each triggered condition.
"""
results = {}
results["confidence_drift"] = self.check_confidence_drift()
results["review_queue_growth"] = self.check_review_queue_growth()
results["low_throughput"] = self.check_low_throughput()
results["high_error_rate"] = self.check_error_rate()
return results
def check_confidence_drift(self) -> bool:
"""
Compare p(C > 0.5) in recent processing logs to the stored baseline.
Returns True if drift exceeds the threshold.
"""
current = self._fetch_recent_confidences(hours=24 * 7)
baseline = self._fetch_baseline_confidences()
if not current or not baseline:
logger.debug("Not enough data for drift check")
return False
triggered = detect_drift(current, baseline, self._drift_threshold)
if triggered:
current_high = _high_conf_fraction(current)
baseline_high = _high_conf_fraction(baseline)
shift = abs(current_high - baseline_high)
self._alert_fn(
f"Confidence drift detected: {shift:.1%} shift in p(C>0.5) "
f"(current={current_high:.1%}, baseline={baseline_high:.1%}, "
f"threshold={self._drift_threshold:.1%})"
)
return triggered
def check_review_queue_growth(self) -> bool:
"""Alert if the unannotated review queue has been above max_size for max_age hours."""
try:
cutoff = datetime.now(timezone.utc) - timedelta(hours=self._review_max_age_h)
rows = self._db.fetchall(
"""SELECT COUNT(*) AS cnt
FROM review_queue
WHERE annotated = FALSE
AND created_at < %s""",
(cutoff.replace(tzinfo=None),),
)
count = int((rows or [{"cnt": 0}])[0]["cnt"])
if count > self._review_max_size:
self._alert_fn(
f"Review queue growth alert: {count} unannotated items older than "
f"{self._review_max_age_h}h (threshold={self._review_max_size})"
)
return True
return False
except Exception as exc:
logger.warning("Review queue growth check failed: %s", exc)
return False
def check_low_throughput(self) -> bool:
"""Alert if throughput dropped below the minimum for the last window hours."""
try:
cutoff = datetime.now(timezone.utc) - timedelta(hours=self._throughput_window_h)
rows = self._db.fetchall(
"""SELECT COUNT(*) AS cnt
FROM videos
WHERE status = 'COMPLETED'
AND last_processed_time >= %s""",
(cutoff.replace(tzinfo=None),),
)
count = int((rows or [{"cnt": 0}])[0]["cnt"])
videos_per_hour = count / max(self._throughput_window_h, 1)
if videos_per_hour < self._throughput_min:
self._alert_fn(
f"Low throughput alert: {videos_per_hour:.1f} videos/hour "
f"over last {self._throughput_window_h}h "
f"(minimum={self._throughput_min})"
)
return True
return False
except Exception as exc:
logger.warning("Throughput check failed: %s", exc)
return False
def check_error_rate(self) -> bool:
"""Alert if the fraction of ERROR/UNSCANNABLE videos exceeds the threshold."""
try:
cutoff = datetime.now(timezone.utc) - timedelta(hours=self._error_rate_window_h)
rows = self._db.fetchall(
"""SELECT status, COUNT(*) AS cnt
FROM videos
WHERE last_processed_time >= %s
AND status IN ('COMPLETED', 'ERROR', 'UNSCANNABLE')
GROUP BY status""",
(cutoff.replace(tzinfo=None),),
)
counts: Dict[str, int] = {r["status"]: int(r["cnt"]) for r in (rows or [])}
total = sum(counts.values())
if total == 0:
return False
errors = counts.get("ERROR", 0) + counts.get("UNSCANNABLE", 0)
rate = errors / total
if rate > self._error_rate_threshold:
self._alert_fn(
f"High error rate alert: {rate:.1%} over last {self._error_rate_window_h}h "
f"({errors}/{total} videos, threshold={self._error_rate_threshold:.1%})"
)
return True
return False
except Exception as exc:
logger.warning("Error rate check failed: %s", exc)
return False
# ------------------------------------------------------------------
# DB helpers
# ------------------------------------------------------------------
def _fetch_recent_confidences(self, hours: int = 168) -> List[float]:
"""Return confidence scores from the last `hours` hours."""
try:
cutoff = datetime.now(timezone.utc) - timedelta(hours=hours)
rows = self._db.fetchall(
"""SELECT confidence_score
FROM processing_logs
WHERE processed_at >= %s
AND confidence_score IS NOT NULL""",
(cutoff.replace(tzinfo=None),),
)
return [float(r["confidence_score"]) for r in (rows or [])]
except Exception as exc:
logger.warning("Failed to fetch recent confidences: %s", exc)
return []
def _fetch_baseline_confidences(self) -> List[float]:
"""Return the baseline confidence distribution stored in the DB."""
try:
rows = self._db.fetchall(
"""SELECT confidence_score
FROM processing_logs
WHERE confidence_score IS NOT NULL
ORDER BY processed_at ASC
LIMIT 10000"""
)
return [float(r["confidence_score"]) for r in (rows or [])]
except Exception as exc:
logger.warning("Failed to fetch baseline confidences: %s", exc)
return []
# ---------------------------------------------------------------------------
# Module-level pure functions (testable without DB)
# ---------------------------------------------------------------------------
def detect_drift(
current_confidences: List[float],
baseline_confidences: List[float],
threshold: float = 0.10,
) -> bool:
"""
Return True if |p(C>0.5)_current - p(C>0.5)_baseline| > threshold.
Both lists must be non-empty.
"""
if not current_confidences or not baseline_confidences:
return False
current_high = _high_conf_fraction(current_confidences)
baseline_high = _high_conf_fraction(baseline_confidences)
shift = abs(current_high - baseline_high)
if shift > threshold:
logger.warning(
"Drift detected: shift=%.2f%% (current=%.2f%% baseline=%.2f%%)",
shift * 100, current_high * 100, baseline_high * 100,
)
return True
return False
def _high_conf_fraction(confidences: List[float]) -> float:
"""Return the fraction of scores > 0.5."""
return sum(1 for c in confidences if c > 0.5) / len(confidences)
def _log_alert(message: str):
logger.warning("ALERT: %s", message)
-104
View File
@@ -1,104 +0,0 @@
"""Export a YOLOv8n face-detection model to ONNX and build a TensorRT engine."""
import argparse
import json
import logging
from pathlib import Path
logger = logging.getLogger(__name__)
def export_onnx(output_dir: Path, input_size: int = 640) -> Path:
"""Export YOLOv8n to ONNX with the required input shape."""
output_dir.mkdir(parents=True, exist_ok=True)
onnx_path = output_dir / "face_detector.onnx"
metadata_path = output_dir / "model.json"
try:
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
model.export(
format="onnx",
imgsz=input_size,
half=False,
simplify=True,
dynamic=False,
)
exported = Path("yolov8n.onnx")
if exported.exists():
exported.rename(onnx_path)
logger.info("Exported ONNX model to %s", onnx_path)
except Exception as exc:
logger.error("Failed to export ONNX model: %s", exc)
raise
metadata = {
"input_shape": [1, 3, input_size, input_size],
"mean": [0.485, 0.456, 0.406],
"std": [0.229, 0.224, 0.225],
"confidence_threshold": 0.25,
"iou_threshold": 0.45,
}
metadata_path.write_text(json.dumps(metadata, indent=2))
return onnx_path
def build_tensorrt_engine(onnx_path: Path, output_dir: Path, max_batch_size: int = 32) -> Path:
"""Build and serialize a TensorRT FP32 engine from an ONNX file."""
engine_path = output_dir / "face_detector.trt"
try:
import tensorrt as trt
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
onnx_data = onnx_path.read_bytes()
if not parser.parse(onnx_data):
for error in range(parser.num_errors):
logger.error(parser.get_error(error))
raise RuntimeError("ONNX parsing failed")
config = builder.create_builder_config()
config.max_workspace_size = 4 << 30 # 4GB
config.set_flag(trt.BuilderFlag.FP32)
profile = builder.create_optimization_profile()
input_name = network.get_input(0).name
profile.set_shape(
input_name,
(1, 3, 640, 640),
(max_batch_size // 2, 3, 640, 640),
(max_batch_size, 3, 640, 640),
)
config.add_optimization_profile(profile)
engine = builder.build_engine(network, config)
if engine is None:
raise RuntimeError("TensorRT engine build failed")
engine_path.write_bytes(engine.serialize())
logger.info("Serialized TensorRT engine to %s", engine_path)
return engine_path
except Exception as exc:
logger.error("Failed to build TensorRT engine: %s", exc)
raise
def main():
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(description="Export YOLOv8n face detector to ONNX/TensorRT")
parser.add_argument("--output-dir", type=Path, default=Path("models/face_detector"))
parser.add_argument("--input-size", type=int, default=640)
parser.add_argument("--max-batch-size", type=int, default=32)
args = parser.parse_args()
onnx_path = export_onnx(args.output_dir, args.input_size)
build_tensorrt_engine(onnx_path, args.output_dir, args.max_batch_size)
if __name__ == "__main__":
main()
-295
View File
@@ -1,295 +0,0 @@
"""
Face detection runner.
Wraps a YOLOv8n model exported to ONNX/TensorRT, runs batched inference on
sampled frames, applies NMS, and produces 224×224 face crops for the classifier.
"""
import logging
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
logger = logging.getLogger(__name__)
class Detection:
"""A single face detection result."""
def __init__(
self,
frame_path: str,
x1: float,
y1: float,
x2: float,
y2: float,
confidence: float,
crop_path: Optional[str] = None,
):
self.frame_path = frame_path
self.x1 = x1
self.y1 = y1
self.x2 = x2
self.y2 = y2
self.confidence = confidence
self.crop_path = crop_path
def to_dict(self) -> Dict[str, Any]:
return {
"frame_path": self.frame_path,
"x1": self.x1,
"y1": self.y1,
"x2": self.x2,
"y2": self.y2,
"confidence": self.confidence,
"crop_path": self.crop_path,
}
class FaceDetector:
"""Run batched face detection on video frames."""
def __init__(
self,
engine_path: str,
input_size: int = 640,
confidence_threshold: float = 0.25,
iou_threshold: float = 0.45,
max_faces_per_frame: int = 10,
max_faces_per_video: int = 100,
device: str = "cuda",
):
self.engine_path = engine_path
self.input_size = input_size
self.confidence_threshold = confidence_threshold
self.iou_threshold = iou_threshold
self.max_faces_per_frame = max_faces_per_frame
self.max_faces_per_video = max_faces_per_video
self.device = device
self._session: Optional[Any] = None
self._load_model()
def _load_model(self):
"""Load the inference backend."""
path = Path(self.engine_path)
if not path.exists():
logger.warning("Face detector engine not found at %s; using placeholder", self.engine_path)
return
suffix = path.suffix.lower()
if suffix == ".onnx":
try:
import onnxruntime as ort
providers = ["CUDAExecutionProvider"] if self.device.startswith("cuda") else ["CPUExecutionProvider"]
self._session = ort.InferenceSession(str(path), providers=providers)
logger.info("Loaded ONNX face detector from %s", self.engine_path)
except Exception as exc:
logger.warning("Failed to load ONNX face detector: %s", exc)
elif suffix in (".trt", ".engine", ".plan"):
try:
import tensorrt as trt
with trt.Logger() as trt_logger, open(path, "rb") as f:
runtime = trt.Runtime(trt_logger)
self._session = runtime.deserialize_cuda_engine(f.read())
logger.info("Loaded TensorRT face detector from %s", self.engine_path)
except Exception as exc:
logger.warning("Failed to load TensorRT face detector: %s", exc)
else:
logger.warning("Unsupported face detector format: %s", suffix)
def detect_faces(
self,
frame_paths: List[str],
batch_size: int = 16,
) -> List[List[Detection]]:
"""Detect faces in a list of frame image paths."""
all_results: List[List[Detection]] = []
for i in range(0, len(frame_paths), batch_size):
batch = frame_paths[i : i + batch_size]
batch_results = self._detect_batch(batch)
all_results.extend(batch_results)
return all_results
def _detect_batch(self, frame_paths: List[str]) -> List[List[Detection]]:
"""Run detection on one batch and return detections per frame."""
preprocessed = []
for path in frame_paths:
try:
preprocessed.append(self._preprocess(path))
except Exception as exc:
logger.warning("Failed to preprocess %s: %s", path, exc)
preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
batch_input = np.stack(preprocessed, axis=0)
outputs = self._infer(batch_input)
results = []
for idx, path in enumerate(frame_paths):
try:
detections = self._parse_outputs(outputs, idx, path)
detections = self._nms(detections)
detections = detections[: self.max_faces_per_frame]
results.append(detections)
except Exception as exc:
logger.warning("Failed to parse detections for %s: %s", path, exc)
results.append([])
return results
def _preprocess(self, frame_path: str) -> np.ndarray:
"""Load and normalize a frame for the detector."""
from PIL import Image
image = Image.open(frame_path).convert("RGB")
image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
arr = np.array(image, dtype=np.float32) / 255.0
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
arr = (arr - mean) / std
return np.transpose(arr, (2, 0, 1)) # HWC -> CHW
def _infer(self, batch_input: np.ndarray) -> Any:
"""Run inference on the preprocessed batch."""
if self._session is None:
return self._placeholder_outputs(batch_input.shape[0])
try:
if hasattr(self._session, "run"):
input_name = self._session.get_inputs()[0].name
return self._session.run(None, {input_name: batch_input})
# TensorRT execution
import pycuda.driver as cuda
import pycuda.autoinit # noqa: F401
context = self._session.create_execution_context()
output_shape = (batch_input.shape[0], 84, 8400) # YOLOv8n default shape
d_input = cuda.mem_alloc(batch_input.nbytes)
d_output = cuda.mem_alloc(np.prod(output_shape) * np.dtype(np.float32).itemsize)
stream = cuda.Stream()
cuda.memcpy_htod_async(d_input, batch_input, stream)
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
output = np.empty(output_shape, dtype=np.float32)
cuda.memcpy_dtoh_async(output, d_output, stream)
stream.synchronize()
return output
except Exception as exc:
logger.warning("Face detector inference failed: %s", exc)
return self._placeholder_outputs(batch_input.shape[0])
def _placeholder_outputs(self, batch_size: int) -> np.ndarray:
"""Return an empty output tensor when no model is loaded."""
return np.zeros((batch_size, 84, 8400), dtype=np.float32)
def _parse_outputs(self, outputs: Any, batch_index: int, frame_path: str) -> List[Detection]:
"""Parse raw inference outputs into Detection objects."""
if isinstance(outputs, list):
raw = outputs[0][batch_index] # (84, 8400)
else:
raw = outputs[batch_index]
# YOLOv8 output layout: (cx, cy, w, h, cls scores...)
raw = raw.T # (8400, 84)
scores = raw[:, 4:].max(axis=1)
mask = scores >= self.confidence_threshold
candidates = raw[mask]
scores = scores[mask]
detections = []
for row, score in zip(candidates, scores):
cx, cy, w, h = row[:4]
x1 = cx - w / 2
y1 = cy - h / 2
x2 = cx + w / 2
y2 = cy + h / 2
detections.append(
Detection(
frame_path=frame_path,
x1=float(x1),
y1=float(y1),
x2=float(x2),
y2=float(y2),
confidence=float(score),
)
)
return detections
def _nms(self, detections: List[Detection]) -> List[Detection]:
"""Apply greedy Non-Maximum Suppression."""
if not detections:
return []
sorted_dets = sorted(detections, key=lambda d: d.confidence, reverse=True)
kept: List[Detection] = []
while sorted_dets:
current = sorted_dets.pop(0)
kept.append(current)
sorted_dets = [
det
for det in sorted_dets
if self._iou(current, det) <= self.iou_threshold
]
if len(kept) >= self.max_faces_per_frame:
break
return kept
@staticmethod
def _iou(a: Detection, b: Detection) -> float:
"""Compute IoU between two detections."""
x1 = max(a.x1, b.x1)
y1 = max(a.y1, b.y1)
x2 = min(a.x2, b.x2)
y2 = min(a.y2, b.y2)
inter_area = max(0, x2 - x1) * max(0, y2 - y1)
area_a = (a.x2 - a.x1) * (a.y2 - a.y1)
area_b = (b.x2 - b.x1) * (b.y2 - b.y1)
union_area = area_a + area_b - inter_area
if union_area == 0:
return 0.0
return inter_area / union_area
def extract_crops(
self,
detections: List[Detection],
output_dir: str,
crop_size: Tuple[int, int] = (224, 224),
) -> List[Detection]:
"""Extract resized face crops from original frames and update detections."""
from PIL import Image
output_dir_path = Path(output_dir)
output_dir_path.mkdir(parents=True, exist_ok=True)
cropped: List[Detection] = []
for idx, det in enumerate(detections):
try:
image = Image.open(det.frame_path).convert("RGB")
width, height = image.size
x1 = int(max(0, det.x1 * width / self.input_size))
y1 = int(max(0, det.y1 * height / self.input_size))
x2 = int(min(width, det.x2 * width / self.input_size))
y2 = int(min(height, det.y2 * height / self.input_size))
if x2 <= x1 or y2 <= y1:
image.close()
continue
crop = image.crop((x1, y1, x2, y2)).resize(crop_size, Image.Resampling.BILINEAR)
image.close()
crop_path = output_dir_path / f"crop_{int(time.time() * 1000)}_{idx}.jpg"
crop.save(crop_path, "JPEG", quality=95)
crop.close()
det.crop_path = str(crop_path)
cropped.append(det)
except Exception as exc:
logger.warning("Failed to extract crop for %s: %s", det.frame_path, exc)
return cropped
-163
View File
@@ -1,163 +0,0 @@
"""
Frame sampling module for VideoDetect.
Extracts uniform temporal frames from videos using FFmpeg, handles resolution
constraints, and writes JPEG frames to a scratch directory.
"""
import logging
import os
import subprocess
from pathlib import Path
from typing import Iterable, List, Optional, Sequence, Tuple
logger = logging.getLogger(__name__)
def calculate_timestamps(duration: Optional[float], interval: float) -> List[float]:
"""Return uniformly spaced timestamps in seconds for a video duration."""
if duration is None or duration <= 0:
return [0.0]
if interval <= 0:
interval = 30.0
count = max(1, int(duration / interval))
step = duration / count
return [round(i * step, 3) for i in range(count)]
class FrameSampler:
"""Extract sampled frames from videos using FFmpeg."""
def __init__(self, interval_seconds: int = 30, quality: int = 2, output_format: str = "jpeg"):
self.interval_seconds = interval_seconds
self.quality = quality
self.output_format = output_format.lower()
def extract_frame(
self,
video_path: str,
output_path: str,
timestamp: float,
resolution: Optional[Tuple[int, int]] = None,
timeout_seconds: int = 30,
) -> Optional[str]:
"""Extract a single frame at the provided timestamp."""
output = Path(output_path)
output.parent.mkdir(parents=True, exist_ok=True)
cmd = self._build_ffmpeg_command(
video_path=video_path,
output_path=str(output),
timestamp=timestamp,
resolution=resolution,
timeout_seconds=timeout_seconds,
)
try:
completed = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=timeout_seconds + 5,
check=False,
)
except subprocess.TimeoutExpired as exc:
logger.warning("Frame extraction timed out for %s at %.3fs", video_path, timestamp)
return None
if completed.returncode != 0:
logger.warning(
"Frame extraction failed for %s at %.3fs: %s",
video_path,
timestamp,
completed.stderr.strip(),
)
return None
if output.exists():
return str(output)
return None
def extract_frames(
self,
video_path: str,
output_dir: str,
duration: Optional[float] = None,
interval_seconds: Optional[int] = None,
resolution: Optional[Tuple[int, int]] = None,
timeout_seconds: int = 30,
) -> List[str]:
"""Extract uniform frames for the video and return output paths."""
resolved_interval = interval_seconds or self.interval_seconds
timestamps = calculate_timestamps(duration=duration, interval=float(resolved_interval))
extracted: List[str] = []
for ts in timestamps:
frame_name = self._frame_name_from_timestamp(ts, video_path)
output_path = str(Path(output_dir) / frame_name)
result = self.extract_frame(
video_path=video_path,
output_path=output_path,
timestamp=ts,
resolution=resolution,
timeout_seconds=timeout_seconds,
)
if result:
extracted.append(result)
return extracted
def _build_ffmpeg_command(
self,
video_path: str,
output_path: str,
timestamp: float,
resolution: Optional[Tuple[int, int]] = None,
timeout_seconds: int = 30,
) -> List[str]:
"""Build an FFmpeg command for extracting a single frame."""
cmd = ["ffmpeg", "-y", "-ss", str(timestamp), "-i", video_path]
vf_filters = []
if resolution is not None:
target_width, target_height = self._resolve_target_resolution(resolution)
vf_filters.append(f"scale={target_width}:{target_height}:force_original_aspect_ratio=decrease")
if vf_filters:
cmd.extend(["-vf", ",".join(vf_filters)])
cmd.extend([
"-vframes",
"1",
"-q:v",
str(self.quality),
"-f",
self.output_format,
output_path,
])
return cmd
@staticmethod
def _resolve_target_resolution(resolution: Tuple[int, int]) -> Tuple[int, int]:
"""Downscale frames larger than 1080p to fit VRAM constraints."""
width, height = resolution
if max(width, height) > 1080:
if width >= 3840 or height >= 2160:
return 1920, 1080
if width >= 2560 or height >= 1440:
return 1280, 720
return width, height
@staticmethod
def _frame_name_from_timestamp(timestamp: float, video_path: str) -> str:
video_id = Path(video_path).stem
timestamp_ms = int(round(timestamp * 1000))
return f"{video_id}_{timestamp_ms}.jpg"
def get_target_resolution(resolution: Optional[Tuple[int, int]]) -> Tuple[int, int]:
"""Compatibility helper for resolution handling."""
if resolution is None:
return (0, 0)
return FrameSampler._resolve_target_resolution(resolution)
-101
View File
@@ -1,101 +0,0 @@
"""GPU memory management and batch-size auto-tuning helpers."""
import logging
from typing import Optional, Tuple
logger = logging.getLogger(__name__)
class GPUMemoryManager:
"""Track GPU memory and recommend safe batch sizes."""
def __init__(
self,
device: int = 0,
max_memory_gb: float = 18.0,
reduce_threshold_gb: float = 16.0,
increase_threshold_gb: float = 10.0,
initial_batch_size: int = 16,
min_batch_size: int = 1,
max_batch_size: int = 32,
):
self.device = device
self.max_memory_gb = max_memory_gb
self.reduce_threshold_gb = reduce_threshold_gb
self.increase_threshold_gb = increase_threshold_gb
self.batch_size = initial_batch_size
self.min_batch_size = min_batch_size
self.max_batch_size = max_batch_size
def get_memory_stats(self) -> Tuple[float, float]:
"""Return (allocated_gb, reserved_gb) for the managed device."""
try:
import torch
if not torch.cuda.is_available():
return 0.0, 0.0
allocated = torch.cuda.memory_allocated(self.device) / (1024**3)
reserved = torch.cuda.memory_reserved(self.device) / (1024**3)
return allocated, reserved
except Exception as exc:
logger.debug("Could not query GPU memory: %s", exc)
return 0.0, 0.0
def adjust_batch_size(self) -> int:
"""Adjust the current batch size based on memory pressure."""
allocated, _ = self.get_memory_stats()
if allocated >= self.reduce_threshold_gb:
new_batch_size = max(self.min_batch_size, int(self.batch_size * 0.75))
if new_batch_size < self.batch_size:
logger.info(
"GPU memory high (%.2f GB), reducing batch size %d -> %d",
allocated,
self.batch_size,
new_batch_size,
)
self.batch_size = new_batch_size
elif allocated <= self.increase_threshold_gb:
new_batch_size = min(self.max_batch_size, int(self.batch_size * 1.25))
if new_batch_size > self.batch_size:
logger.info(
"GPU memory low (%.2f GB), increasing batch size %d -> %d",
allocated,
self.batch_size,
new_batch_size,
)
self.batch_size = new_batch_size
return self.batch_size
def is_memory_critical(self) -> bool:
"""Return True if allocated memory is close to the hard limit."""
allocated, _ = self.get_memory_stats()
return allocated >= self.max_memory_gb
def empty_cache(self):
"""Try to release unused cached GPU memory."""
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception as exc:
logger.debug("Could not empty GPU cache: %s", exc)
@property
def current_batch_size(self) -> int:
return self.batch_size
def get_available_vram_gb(device: int = 0) -> float:
"""Return total available VRAM in GB, or 0.0 if CUDA is unavailable."""
try:
import torch
if not torch.cuda.is_available():
return 0.0
return torch.cuda.get_device_properties(device).total_memory / (1024**3)
except Exception:
return 0.0
-134
View File
@@ -1,134 +0,0 @@
"""
Worker health check HTTP endpoint.
Exposes GET /health returning a JSON status document.
Designed to be started in a background thread alongside the main worker.
"""
import json
import logging
import os
import threading
import time
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Any, Dict, Optional
logger = logging.getLogger(__name__)
_START_TIME = time.monotonic()
class HealthStatus:
"""Mutable singleton updated by the worker as it processes videos."""
def __init__(self):
self._lock = threading.Lock()
self._data: Dict[str, Any] = {
"status": "starting",
"gpu_available": False,
"gpu_memory_used_gb": 0.0,
"queue_depth": 0,
"videos_processed_today": 0,
"last_error": None,
}
def update(self, **kwargs):
with self._lock:
self._data.update(kwargs)
def snapshot(self) -> Dict[str, Any]:
with self._lock:
data = dict(self._data)
data["uptime_seconds"] = round(time.monotonic() - _START_TIME, 1)
return data
# Module-level singleton shared between the HTTP handler and the worker
health = HealthStatus()
class _HealthHandler(BaseHTTPRequestHandler):
"""Minimal HTTP handler — only responds to GET /health."""
def do_GET(self): # noqa: N802
if self.path.rstrip("/") == "/health":
body = json.dumps(health.snapshot(), default=str).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
else:
self.send_response(404)
self.end_headers()
def log_message(self, fmt, *args): # silence default access log
logger.debug("health_check: " + fmt, *args)
def start_health_server(port: int = 8080) -> threading.Thread:
"""
Start the health HTTP server in a daemon thread.
Returns the thread so callers can join if needed.
"""
server = HTTPServer(("0.0.0.0", port), _HealthHandler)
def _serve():
logger.info("Health check server listening on port %d", port)
server.serve_forever()
t = threading.Thread(target=_serve, daemon=True, name="health-server")
t.start()
return t
def refresh_health_from_system(db_connector=None, scratch_path: str = "/scratch"):
"""
Update the health singleton with current GPU, queue, and disk state.
Call this periodically from the worker main loop.
"""
updates: Dict[str, Any] = {"status": "healthy"}
# GPU state
try:
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
updates["gpu_available"] = True
updates["gpu_memory_used_gb"] = round(mem.used / 1024**3, 2)
pynvml.nvmlShutdown()
except Exception:
updates["gpu_available"] = _torch_gpu_available()
# Queue depth from DB
if db_connector is not None:
try:
rows = db_connector.fetchall(
"SELECT COUNT(*) AS cnt FROM videos WHERE status IN ('PENDING','PROCESSING')"
)
updates["queue_depth"] = int((rows or [{"cnt": 0}])[0]["cnt"])
today = datetime.now(timezone.utc).date().isoformat()
rows2 = db_connector.fetchall(
"SELECT COUNT(*) AS cnt FROM videos "
"WHERE status = 'COMPLETED' AND DATE(last_processed_time) = %s",
(today,),
)
updates["videos_processed_today"] = int((rows2 or [{"cnt": 0}])[0]["cnt"])
except Exception as exc:
logger.debug("Health DB query failed: %s", exc)
health.update(**updates)
def _torch_gpu_available() -> bool:
try:
import torch
return torch.cuda.is_available()
except ImportError:
return False
-87
View File
@@ -1,87 +0,0 @@
"""
Logging configuration for VideoDetect.
Sets up JSON structured logging via python-json-logger with:
- Configurable log levels
- Standardized field names
- Log rotation
- All logs to stdout for Docker capture
"""
import json
import logging
import logging.handlers
import sys
from datetime import datetime, timezone
def setup_logging(
level: str = "INFO",
log_format: str = "json",
rotation_max_bytes: int = 104857600, # 100MB
rotation_backup_count: int = 10,
):
"""Configure structured logging for the application."""
log_level = getattr(logging, level.upper(), logging.INFO)
# Root logger
root_logger = logging.getLogger()
root_logger.setLevel(log_level)
# Remove existing handlers
root_logger.handlers.clear()
# Handler: stdout (for Docker capture)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(log_level)
if log_format == "json":
formatter = JsonFormatter()
else:
formatter = logging.Formatter(
"%(asctime)s %(levelname)s %(name)s %(message)s",
datefmt="%Y-%m-%dT%H:%M:%S%z",
)
handler.setFormatter(formatter)
root_logger.addHandler(handler)
# Handler: file rotation (for persistence)
file_handler = logging.handlers.RotatingFileHandler(
filename="/logs/videodetect.log",
maxBytes=rotation_max_bytes,
backupCount=rotation_backup_count,
encoding="utf-8",
)
file_handler.setLevel(log_level)
file_handler.setFormatter(formatter)
root_logger.addHandler(file_handler)
logging.info("Logging configured: level=%s format=%s", level, log_format)
class JsonFormatter(logging.Formatter):
"""JSON log formatter with standardized fields."""
def format(self, record: logging.LogRecord) -> str:
log_data = {
"timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
"level": record.levelname,
"logger": record.name,
"message": record.getMessage(),
"module": record.module,
"function": record.funcName,
"line": record.lineno,
}
# Add extra fields
if hasattr(record, "video_id"):
log_data["video_id"] = record.video_id
if hasattr(record, "metadata"):
log_data["metadata"] = record.metadata
# Add exception info if present
if record.exc_info and record.exc_info[0] is not None:
log_data["exception"] = self.formatException(record.exc_info)
return json.dumps(log_data, default=str)
-122
View File
@@ -1,122 +0,0 @@
"""
VideoDetect - Video Classification System
Main entry point for the worker service.
Initializes all components and starts the processing pipeline.
"""
import logging
import signal
import sys
import threading
import time
from pathlib import Path
# Add src to path
sys.path.insert(0, str(Path(__file__).parent))
from config_loader import get_config
from db_connector import DBConnector
from logging_config import setup_logging
from orchestrator import WorkerPool
from scanner import DirectoryScanner
logger = logging.getLogger(__name__)
def graceful_shutdown(signum, frame):
"""Handle shutdown signals gracefully."""
logger.info("Received signal %d, initiating graceful shutdown...", signum)
sys.exit(0)
def main():
"""Initialize and start the VideoDetect worker."""
# Register signal handlers
signal.signal(signal.SIGTERM, graceful_shutdown)
signal.signal(signal.SIGINT, graceful_shutdown)
# Load configuration
config = get_config()
log_config = config.get_section("logging")
setup_logging(
level=log_config.get("level", "INFO"),
log_format=log_config.get("format", "json"),
rotation_max_bytes=log_config.get("rotation_max_bytes", 104857600),
rotation_backup_count=log_config.get("rotation_backup_count", 10),
)
logger.info("VideoDetect Worker starting...")
logger.info("Config: %s", config)
# Initialize database connection
db_config = config.get_section("database")
db = DBConnector(
host=db_config.get("host", "mariadb"),
port=db_config.get("port", 3306),
database=db_config.get("name", "videodetect"),
user=db_config.get("user", "videodetect"),
password=db_config.get("password", "videodetect123"),
pool_size=db_config.get("pool_size", 20),
pool_min=db_config.get("pool_min", 5),
pool_recycle=db_config.get("pool_recycle", 3600),
)
# Verify database connectivity
if not db.health_check():
logger.error("Cannot connect to database. Exiting.")
sys.exit(1)
logger.info("Database connection established.")
# Initialize schema if needed
schema_path = Path(__file__).parent.parent / "db" / "schema.sql"
if schema_path.exists():
db.initialize_schema(str(schema_path))
logger.info("Schema initialized.")
# Verify GPU availability
try:
import torch
if torch.cuda.is_available():
gpu_count = torch.cuda.device_count()
gpu_name = torch.cuda.get_device_name(0)
logger.info("GPU available: %d GPUs, primary: %s", gpu_count, gpu_name)
else:
logger.warning("CUDA is not available! Processing will be slow.")
except ImportError:
logger.warning("PyTorch not installed. GPU features disabled.")
# Initialize the directory scanner.
# It scans the permanent storage location in place (no staging/copy step),
# periodically discovering new/removed video files and queueing any that
# have not been scanned yet as PENDING for the worker pool to pick up.
storage_config = config.get_section("storage")
scanner_config = config.get_section("scanner")
scanner = DirectoryScanner(
db_connector=db,
config=config.data,
input_path=storage_config.get("input_path", "/data/input"),
scan_interval=int(scanner_config.get("scan_interval_seconds", 60)),
walker_threads=int(scanner_config.get("walker_threads", 8)),
)
logger.info("Worker initialization complete. Starting scanner and processing loop...")
# Run the scanner in a background thread (its start() is a blocking loop).
scanner_thread = threading.Thread(
target=scanner.start, name="directory-scanner", daemon=True
)
scanner_thread.start()
pool = WorkerPool(db, config.data, max_workers=1)
try:
pool.start()
except KeyboardInterrupt:
logger.info("Worker shutting down.")
pool.stop()
scanner.stop()
if __name__ == "__main__":
main()
-261
View File
@@ -1,261 +0,0 @@
"""
Prometheus metrics definitions for VideoDetect.
All metric objects are module-level singletons — import and call them directly.
Call start_metrics_server() once at worker startup to expose /metrics.
"""
import logging
import threading
from typing import Optional
logger = logging.getLogger(__name__)
# Populated on first _ensure_prometheus() call; None if library unavailable.
_prometheus_available: Optional[bool] = None
_server_started = False
_server_lock = threading.Lock()
# ---------------------------------------------------------------------------
# Lazy metric definitions — avoid import errors when prometheus_client absent
# ---------------------------------------------------------------------------
def _ensure_prometheus():
global _prometheus_available
if _prometheus_available is not None:
return _prometheus_available
try:
import prometheus_client # noqa: F401
_prometheus_available = True
except ImportError:
logger.warning("prometheus_client not installed; metrics disabled")
_prometheus_available = False
return _prometheus_available
class _NoOpMetric:
"""Drop-in replacement used when prometheus_client is unavailable."""
def labels(self, **_):
return self
def inc(self, *_, **__):
pass
def set(self, *_, **__):
pass
def observe(self, *_, **__):
pass
def _counter(name, doc, labelnames=()):
if not _ensure_prometheus():
return _NoOpMetric()
from prometheus_client import Counter
return Counter(name, doc, list(labelnames))
def _gauge(name, doc, labelnames=()):
if not _ensure_prometheus():
return _NoOpMetric()
from prometheus_client import Gauge
return Gauge(name, doc, list(labelnames))
def _histogram(name, doc, labelnames=(), buckets=None):
if not _ensure_prometheus():
return _NoOpMetric()
from prometheus_client import Histogram
kwargs = {"labelnames": list(labelnames)}
if buckets is not None:
kwargs["buckets"] = buckets
return Histogram(name, doc, **kwargs)
# ---------------------------------------------------------------------------
# Counters
# ---------------------------------------------------------------------------
videos_processed_total = _counter(
"videos_processed_total",
"Total videos processed",
["routing_decision"],
)
videos_processed_by_model_total = _counter(
"videos_processed_by_model_total",
"Total videos processed per model version",
["model_version"],
)
frames_extracted_total = _counter(
"frames_extracted_total",
"Total frames extracted across all videos",
)
faces_detected_total = _counter(
"faces_detected_total",
"Total face crops detected",
)
inference_errors_total = _counter(
"inference_errors_total",
"Total inference errors by type",
["error_type"],
)
retry_attempts_total = _counter(
"retry_attempts_total",
"Total retry attempts by pipeline step",
["step"],
)
# ---------------------------------------------------------------------------
# Gauges
# ---------------------------------------------------------------------------
gpu_utilization_percent = _gauge(
"gpu_utilization_percent",
"GPU utilization as a percentage",
["gpu_id"],
)
gpu_memory_used_bytes = _gauge(
"gpu_memory_used_bytes",
"GPU memory currently allocated in bytes",
["gpu_id"],
)
gpu_memory_free_bytes = _gauge(
"gpu_memory_free_bytes",
"GPU memory currently free in bytes",
["gpu_id"],
)
queue_depth_pending = _gauge(
"queue_depth_pending",
"Number of videos in PENDING state",
)
queue_depth_processing = _gauge(
"queue_depth_processing",
"Number of videos currently in PROCESSING state",
)
queue_depth_review = _gauge(
"queue_depth_review",
"Number of unannotated videos in the review queue",
)
scratch_usage_bytes = _gauge(
"scratch_usage_bytes",
"Scratch space used in bytes",
)
scratch_usage_percent = _gauge(
"scratch_usage_percent",
"Scratch space used as a percentage of total",
)
# ---------------------------------------------------------------------------
# Histograms
# ---------------------------------------------------------------------------
_CONFIDENCE_BUCKETS = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
video_processing_duration_seconds = _histogram(
"video_processing_duration_seconds",
"End-to-end processing time per video in seconds",
["routing_decision"],
)
confidence_score_distribution = _histogram(
"confidence_score_distribution",
"Distribution of video-level confidence scores",
buckets=_CONFIDENCE_BUCKETS,
)
frame_count_per_video = _histogram(
"frame_count_per_video",
"Number of frames sampled per video",
buckets=[1, 2, 5, 10, 20, 50, 100, 200, 500],
)
face_count_per_video = _histogram(
"face_count_per_video",
"Number of face crops detected per video",
buckets=[0, 1, 2, 5, 10, 20, 50, 100],
)
# ---------------------------------------------------------------------------
# Server startup
# ---------------------------------------------------------------------------
def start_metrics_server(port: int = 9090):
"""Start the Prometheus HTTP metrics endpoint (idempotent)."""
global _server_started
with _server_lock:
if _server_started:
return
if not _ensure_prometheus():
logger.warning("Cannot start metrics server: prometheus_client unavailable")
return
try:
from prometheus_client import start_http_server
start_http_server(port)
_server_started = True
logger.info("Prometheus metrics server started on port %d", port)
except Exception as exc:
logger.error("Failed to start metrics server: %s", exc)
# ---------------------------------------------------------------------------
# Convenience update functions
# ---------------------------------------------------------------------------
def update_queue_depths(db_connector):
"""Query the DB and refresh all queue-depth gauges."""
try:
rows = db_connector.fetchall(
"""SELECT status, COUNT(*) AS cnt
FROM videos
WHERE status IN ('PENDING', 'PROCESSING')
GROUP BY status"""
)
counts = {r["status"]: int(r["cnt"]) for r in (rows or [])}
queue_depth_pending.set(counts.get("PENDING", 0))
queue_depth_processing.set(counts.get("PROCESSING", 0))
review_rows = db_connector.fetchall(
"SELECT COUNT(*) AS cnt FROM review_queue WHERE annotated = FALSE"
)
queue_depth_review.set(int((review_rows or [{"cnt": 0}])[0]["cnt"]))
except Exception as exc:
logger.warning("Failed to update queue depth metrics: %s", exc)
def update_gpu_metrics():
"""Poll nvidia-smi (via pynvml) and refresh GPU gauges."""
try:
import pynvml
pynvml.nvmlInit()
count = pynvml.nvmlDeviceGetCount()
for i in range(count):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
gpu_id = str(i)
gpu_utilization_percent.labels(gpu_id=gpu_id).set(util.gpu)
gpu_memory_used_bytes.labels(gpu_id=gpu_id).set(mem.used)
gpu_memory_free_bytes.labels(gpu_id=gpu_id).set(mem.free)
pynvml.nvmlShutdown()
except Exception as exc:
logger.debug("GPU metrics unavailable: %s", exc)
def update_scratch_metrics(scratch_path: str = "/scratch"):
"""Measure scratch disk usage and refresh gauges."""
try:
import shutil
total, used, free = shutil.disk_usage(scratch_path)
scratch_usage_bytes.set(used)
scratch_usage_percent.set(round(used / max(total, 1) * 100, 2))
except Exception as exc:
logger.debug("Scratch metrics unavailable: %s", exc)
-441
View File
@@ -1,441 +0,0 @@
"""
Batch Orchestration Skeleton
Manages job queue, worker pool, state transitions, and crash recovery.
Ensures atomic state transitions and idempotent processing.
"""
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
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
logger = logging.getLogger(__name__)
class JobStatus(Enum):
"""Job status states."""
PENDING = "PENDING"
PROCESSING = "PROCESSING"
COMPLETED = "COMPLETED"
FAILED = "ERROR"
SKIPPED = "UNSCANNABLE"
class Job:
"""Represents a single video processing job."""
def __init__(
self,
video_id: int,
file_path: str,
priority: float = 0.0,
sampling_interval_seconds: Optional[int] = None,
):
self.video_id = video_id
self.file_path = file_path
self.priority = priority # Higher = more urgent (based on modification time)
self.sampling_interval_seconds = sampling_interval_seconds
self.status = JobStatus.PENDING
self.created_at = datetime.now(timezone.utc)
self.started_at: Optional[datetime] = None
self.completed_at: Optional[datetime] = None
self.attempts = 0
self.max_retries = 3
self.error_message: Optional[str] = None
def to_dict(self) -> dict:
return {
"video_id": self.video_id,
"file_path": self.file_path,
"priority": self.priority,
"status": self.status.value,
"created_at": self.created_at.isoformat(),
"started_at": self.started_at.isoformat() if self.started_at else None,
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
"attempts": self.attempts,
"error_message": self.error_message,
}
def __repr__(self):
return f"Job(id={self.video_id}, path={self.file_path}, status={self.status.value})"
class WorkerPool:
"""Manage a pool of worker processes for parallel processing."""
def __init__(self, db_connector, config: dict, max_workers: int = 2):
self.db = db_connector
self.config = config
self.max_workers = max_workers
self._executor = ThreadPoolExecutor(max_workers=max_workers)
self._running = False
self._jobs_processed = 0
self._jobs_failed = 0
self._sampling_config = (config or {}).get("sampling", {})
self._storage_config = (config or {}).get("storage", {})
self._face_detection_config = (config or {}).get("face_detection", {})
self._batching_config = (config or {}).get("batching", {})
self._gpu_manager = GPUMemoryManager(
max_memory_gb=config.get("gpu", {}).get("max_memory_gb", 18.0),
reduce_threshold_gb=self._batching_config.get("vram_reduce_threshold_gb", 16.0),
increase_threshold_gb=self._batching_config.get("vram_increase_threshold_gb", 10.0),
initial_batch_size=self._batching_config.get("max_batch_size", 16),
)
self._face_detector = FaceDetector(
engine_path=self._face_detection_config.get("model_path", "/models/face_detector/face_detector.trt"),
input_size=int(self._face_detection_config.get("input_size", 640)),
confidence_threshold=float(self._face_detection_config.get("confidence_threshold", 0.25)),
iou_threshold=float(self._face_detection_config.get("iou_threshold", 0.45)),
max_faces_per_frame=int(self._face_detection_config.get("max_faces_per_frame", 10)),
max_faces_per_video=int(self._face_detection_config.get("max_faces_per_video", 100)),
)
classifier_config = (config or {}).get("classifier", {})
self._classifier = FaceClassifier(
engine_path=classifier_config.get("model_path", "/models/classifier/classifier.trt"),
temperature=float(classifier_config.get("temperature", 1.0)),
input_size=int(classifier_config.get("input_size", 224)),
)
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."""
self._running = True
logger.info("Worker pool starting with %d workers", self.max_workers)
while self._running:
# Get pending jobs
jobs = self._get_pending_jobs()
if jobs:
# Submit jobs to executor
for job in jobs:
future = self._executor.submit(self._process_job, job)
future.add_done_callback(self._on_job_complete)
else:
time.sleep(5) # No jobs, wait
def stop(self):
"""Stop the worker pool gracefully."""
self._running = False
logger.info("Worker pool stopping. Processed: %d, Failed: %d",
self._jobs_processed, self._jobs_failed)
self._executor.shutdown(wait=True)
def _get_pending_jobs(self) -> List[Job]:
"""Get pending jobs from DB with atomic locking."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
# Atomic lock: update status from PENDING to PROCESSING
result = self.db.execute(
"""UPDATE videos SET status = 'PROCESSING', last_processed_time = %s,
updated_at = %s
WHERE id IN (
SELECT id FROM (
SELECT id FROM videos
WHERE status = 'PENDING'
ORDER BY last_scan_time DESC
LIMIT %s
) AS sub
)""",
(now, now, self.max_workers),
transaction=True,
)
if result == 0:
return []
# Get the locked jobs
locked = self.db.fetchall(
"SELECT id, file_path, last_scan_time FROM videos WHERE status = 'PROCESSING' AND last_processed_time = %s",
(now,),
)
return [
Job(
video_id=row["id"],
file_path=row["file_path"],
priority=row["last_scan_time"].timestamp(),
)
for row in locked
]
def _process_job(self, job: Job) -> bool:
"""Process a single job (placeholder - actual processing in later stories)."""
job.status = JobStatus.PROCESSING
job.started_at = datetime.now(timezone.utc)
job.attempts += 1
logger.info("Processing job: %s (attempt %d)", job, job.attempts)
try:
sampling_interval = job.sampling_interval_seconds or self._sampling_config.get("interval_seconds", 30)
scratch_base_path = self._storage_config.get("scratch_path", "/scratch")
scratch_manager = ScratchManager(
base_path=scratch_base_path,
video_id=str(job.video_id),
auto_cleanup=True,
)
frame_dir = scratch_manager.ensure_frame_dir()
prober = VideoProber(timeout=10)
metadata = prober.probe(job.file_path)
if metadata.is_unscannable or metadata.duration is None:
raise RuntimeError(metadata.error_message or "Video metadata could not be determined")
resolution = None
if metadata.resolution_w and metadata.resolution_h:
resolution = (metadata.resolution_w, metadata.resolution_h)
sampler = FrameSampler(
interval_seconds=int(sampling_interval),
quality=int(self._sampling_config.get("quality", 2)),
output_format=self._sampling_config.get("format", "jpeg"),
)
extracted_frames = sampler.extract_frames(
video_path=job.file_path,
output_dir=str(frame_dir),
duration=metadata.duration,
interval_seconds=int(sampling_interval),
resolution=resolution,
timeout_seconds=30,
)
if not extracted_frames:
scratch_manager.cleanup()
raise RuntimeError("No frames were extracted")
# Face detection on sampled frames
batch_size = self._gpu_manager.current_batch_size
detections_per_frame = self._face_detector.detect_faces(extracted_frames, batch_size=batch_size)
# Flatten and cap total faces per video
all_detections: List[Detection] = []
for frame_dets in detections_per_frame:
all_detections.extend(frame_dets)
all_detections = sorted(all_detections, key=lambda d: d.confidence, reverse=True)
all_detections = all_detections[: self._face_detection_config.get("max_faces_per_video", 100)]
# Extract face crops
crop_dir = scratch_manager.frame_dir.parent / "crops"
cropped_detections = self._face_detector.extract_crops(
all_detections,
output_dir=str(crop_dir),
crop_size=(224, 224),
)
if not cropped_detections:
logger.info("No faces detected for video %s; routing to SKIP", job.video_id)
routing_decision = router.SKIP
video_confidence = 0.0
frame_confidences: List[float] = []
else:
crop_paths = [d.crop_path for d in cropped_detections if d.crop_path]
frame_confidences = self._classifier.classify(
crop_paths, batch_size=self._gpu_manager.current_batch_size
)
video_confidence = aggregator.aggregate(
frame_confidences,
strategy=self._aggregation_config.get("strategy", "max"),
alpha=float(self._aggregation_config.get("alpha", 1.0)),
beta=float(self._aggregation_config.get("beta", 0.1)),
top_k=int(self._aggregation_config.get("top_k", 3)),
)
routing_decision = router.route(
video_confidence,
t_high=float(self._routing_config.get("T_high", 0.75)),
t_low=float(self._routing_config.get("T_low", 0.45)),
)
logger.info(
"Video %s: C=%.4f routing=%s faces=%d frames=%d",
job.video_id, video_confidence, routing_decision,
len(cropped_detections), len(extracted_frames),
)
# 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,
)
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
except Exception as e:
logger.error("Job failed: %s - %s", job, e, exc_info=True)
job.error_message = str(e)
if job.attempts < job.max_retries:
# Retry
self._retry_job(job)
return False
else:
# Max retries exceeded
self._fail_job(job)
self._jobs_failed += 1
return False
def _complete_job(
self,
job: Job,
frame_count: Optional[int] = None,
confidence: Optional[float] = None,
routing: Optional[str] = None,
):
"""Mark a job as completed, persisting confidence and routing decision."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
job.status = JobStatus.COMPLETED
job.completed_at = datetime.now(timezone.utc)
self.db.execute(
"""UPDATE videos
SET status = 'COMPLETED',
frame_count = %s,
confidence_score = %s,
routing_decision = %s,
updated_at = %s
WHERE id = %s""",
(frame_count, confidence, routing, now, job.video_id),
transaction=True,
)
logger.info("Job completed: %s", job)
def _fail_job(self, job: Job):
"""Mark a job as failed."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
job.status = JobStatus.FAILED
job.completed_at = datetime.now(timezone.utc)
self.db.execute(
"""UPDATE videos SET status = 'ERROR', error_message = %s, updated_at = %s
WHERE id = %s""",
(job.error_message, now, job.video_id),
transaction=True,
)
logger.error("Job failed permanently: %s", job)
def _retry_job(self, job: Job):
"""Re-queue a job for retry."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
self.db.execute(
"""UPDATE videos SET status = 'PENDING', updated_at = %s
WHERE id = %s""",
(now, job.video_id),
transaction=True,
)
logger.info("Job re-queued for retry: %s (attempt %d/%d)",
job, job.attempts, job.max_retries)
def _on_job_complete(self, future):
"""Callback when a job completes."""
try:
future.result()
except Exception as e:
logger.error("Unhandled job error: %s", e)
def get_stats(self) -> dict:
"""Get worker pool statistics."""
return {
"max_workers": self.max_workers,
"jobs_processed": self._jobs_processed,
"jobs_failed": self._jobs_failed,
"is_running": self._running,
"active_workers": self._executor._work_queue.qsize(),
}
class CrashRecovery:
"""Handle crash recovery and stale job detection."""
def __init__(self, db_connector, lock_timeout_minutes: int = 5):
self.db = db_connector
self.lock_timeout = lock_timeout_minutes
def recover_stale_jobs(self) -> int:
"""Re-queue jobs stuck in PROCESSING beyond the lock timeout."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
timeout = now.replace(minute=now.minute - self.lock_timeout)
result = self.db.execute(
"""UPDATE videos SET status = 'PENDING', updated_at = %s
WHERE status = 'PROCESSING' AND last_processed_time < %s""",
(now, timeout),
transaction=True,
)
if result > 0:
logger.info("Recovered %d stale jobs", result)
return result
def check_health(self) -> dict:
"""Check system health for crash recovery purposes."""
processing_count = self.db.fetchone(
"SELECT COUNT(*) as count FROM videos WHERE status = 'PROCESSING'"
)
pending_count = self.db.fetchone(
"SELECT COUNT(*) as count FROM videos WHERE status = 'PENDING'"
)
return {
"processing_count": processing_count["count"] if processing_count else 0,
"pending_count": pending_count["count"] if pending_count else 0,
"stale_jobs_recovered": self.recover_stale_jobs(),
}
-182
View File
@@ -1,182 +0,0 @@
"""
Video Probing Module
Extracts video metadata (codec, resolution, duration, frame rate) using ffprobe.
Handles errors gracefully for corrupt or unsupported files.
"""
import json
import logging
import os
import subprocess
import time
from dataclasses import dataclass, field
from typing import Optional
logger = logging.getLogger(__name__)
@dataclass
class VideoMetadata:
"""Metadata extracted from a video file."""
file_path: str
codec: Optional[str] = None
resolution_w: Optional[int] = None
resolution_h: Optional[int] = None
duration: Optional[float] = None
frame_rate: Optional[float] = None
file_size: Optional[int] = None
status: str = "PENDING"
error_message: Optional[str] = None
probe_time: float = field(default_factory=time.time)
@property
def is_valid(self) -> bool:
return self.status == "PENDING" and self.codec is not None
@property
def is_unscannable(self) -> bool:
return self.status == "UNSCANNABLE"
def to_dict(self) -> dict:
return {
"file_path": self.file_path,
"codec": self.codec,
"resolution_w": self.resolution_w,
"resolution_h": self.resolution_h,
"duration": self.duration,
"frame_rate": self.frame_rate,
"file_size": self.file_size,
"status": self.status,
"error_message": self.error_message,
}
class VideoProber:
"""Probe video files using ffprobe to extract metadata."""
def __init__(self, timeout: int = 10):
self.timeout = timeout
def probe(self, file_path: str) -> VideoMetadata:
"""Probe a video file and return metadata."""
metadata = VideoMetadata(file_path=file_path)
# Get file size
try:
metadata.file_size = self._get_file_size(file_path)
except OSError as e:
metadata.status = "UNSCANNABLE"
metadata.error_message = f"Cannot read file: {e}"
logger.warning("File access error for %s: %s", file_path, e)
return metadata
# Run ffprobe
try:
result = self._run_ffprobe(file_path)
if result is None:
metadata.status = "UNSCANNABLE"
metadata.error_message = "ffprobe timed out or failed"
logger.warning("ffprobe failed for %s", file_path)
return metadata
metadata = self._parse_ffprobe_output(result, file_path)
except subprocess.TimeoutExpired:
metadata.status = "UNSCANNABLE"
metadata.error_message = f"ffprobe timed out after {self.timeout}s"
logger.warning("ffprobe timeout for %s", file_path)
except Exception as e:
metadata.status = "UNSCANNABLE"
metadata.error_message = f"ffprobe error: {e}"
logger.warning("ffprobe error for %s: %s", file_path, e)
return metadata
def _run_ffprobe(self, file_path: str) -> Optional[str]:
"""Run ffprobe and return JSON output."""
cmd = [
"ffprobe",
"-v", "error",
"-select_streams", "v:0",
"-show_entries", "stream=codec_name,width,height,r_frame_rate,duration",
"-show_entries", "format=duration,size",
"-of", "json",
file_path,
]
try:
output = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=self.timeout,
)
if output.returncode != 0:
logger.debug("ffprobe stderr for %s: %s", file_path, output.stderr)
return None
return output.stdout
except subprocess.TimeoutExpired:
return None
def _parse_ffprobe_output(self, output: str, file_path: str) -> VideoMetadata:
"""Parse ffprobe JSON output into VideoMetadata."""
metadata = VideoMetadata(file_path=file_path)
try:
data = json.loads(output)
except json.JSONDecodeError as e:
metadata.status = "UNSCANNABLE"
metadata.error_message = f"Invalid ffprobe output: {e}"
return metadata
# Extract stream info
streams = data.get("streams", [])
if not streams:
metadata.status = "UNSCANNABLE"
metadata.error_message = "No video streams found"
return metadata
stream = streams[0]
metadata.codec = stream.get("codec_name")
metadata.resolution_w = stream.get("width")
metadata.resolution_h = stream.get("height")
# Parse frame rate
r_frame_rate = stream.get("r_frame_rate")
if r_frame_rate and "/" in r_frame_rate:
try:
num, den = r_frame_rate.split("/")
metadata.frame_rate = float(num) / float(den) if float(den) != 0 else None
except (ValueError, ZeroDivisionError):
pass
# Extract format info
format_info = data.get("format", {})
if format_info.get("duration"):
try:
metadata.duration = float(format_info["duration"])
except (ValueError, TypeError):
pass
if format_info.get("size"):
try:
metadata.file_size = int(format_info["size"])
except (ValueError, TypeError):
pass
# Validate we got essential data
if not metadata.codec:
metadata.status = "UNSCANNABLE"
metadata.error_message = "No codec information available"
elif not metadata.duration:
metadata.status = "UNSCANNABLE"
metadata.error_message = "No duration information available"
return metadata
@staticmethod
def _get_file_size(file_path: str) -> int:
"""Get file size in bytes."""
return os.path.getsize(file_path)
-39
View File
@@ -1,39 +0,0 @@
"""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)
-75
View File
@@ -1,75 +0,0 @@
"""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
-114
View File
@@ -1,114 +0,0 @@
"""
Retry decorator with exponential backoff for transient failures.
Usage:
@retry(max_attempts=3, step="extract")
def extract_frames(...): ...
Non-retryable error types are re-raised immediately without consuming attempts.
"""
import functools
import logging
import time
from typing import Callable, Optional, Tuple, Type
logger = logging.getLogger(__name__)
# Error class names that should never be retried
_NON_RETRYABLE_NAMES = frozenset({
"CodecUnsupportedError",
"FileCorruptError",
"InvalidPathError",
"FileNotFoundError",
"PermissionError",
"IsADirectoryError",
})
class RetryExhaustedError(Exception):
"""Raised when all retry attempts have been exhausted."""
def retry(
max_attempts: int = 3,
initial_delay: float = 1.0,
backoff_factor: float = 2.0,
exceptions: Tuple[Type[Exception], ...] = (Exception,),
step: str = "unknown",
):
"""
Decorator: retry ``func`` up to ``max_attempts`` times on retryable exceptions.
Non-retryable exceptions (see _NON_RETRYABLE_NAMES) propagate immediately.
Each retry waits ``initial_delay * backoff_factor ** attempt`` seconds.
"""
def decorator(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(*args, **kwargs):
last_exc: Optional[Exception] = None
for attempt in range(max_attempts):
try:
return func(*args, **kwargs)
except exceptions as exc:
if _is_non_retryable(exc):
logger.error(
"Non-retryable error in step=%s (attempt %d/%d): %s: %s",
step, attempt + 1, max_attempts,
type(exc).__name__, exc,
)
raise
last_exc = exc
_increment_retry_counter(step)
if attempt < max_attempts - 1:
delay = initial_delay * (backoff_factor ** attempt)
logger.warning(
"Transient error in step=%s (attempt %d/%d), "
"retrying in %.1fs: %s: %s",
step, attempt + 1, max_attempts, delay,
type(exc).__name__, exc,
)
time.sleep(delay)
else:
logger.error(
"All %d attempts exhausted in step=%s: %s: %s",
max_attempts, step, type(exc).__name__, exc,
)
raise RetryExhaustedError(
f"step={step} failed after {max_attempts} attempts"
) from last_exc
return wrapper
return decorator
def retry_from_config(config, step: str = "unknown"):
"""Build a ``@retry`` decorator from config.yaml monitoring.retry settings."""
mon = config.get_section("monitoring")
cfg = mon.get("retry", {})
return retry(
max_attempts=int(cfg.get("max_attempts", 3)),
initial_delay=float(cfg.get("initial_delay", 1.0)),
backoff_factor=float(cfg.get("backoff_factor", 2.0)),
step=step,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _is_non_retryable(exc: Exception) -> bool:
return type(exc).__name__ in _NON_RETRYABLE_NAMES
def _increment_retry_counter(step: str):
try:
from metrics import retry_attempts_total
retry_attempts_total.labels(step=step).inc()
except Exception:
pass # metrics not available; don't break the retry logic
-114
View File
@@ -1,114 +0,0 @@
"""
Review queue export for backend scripts.
Provides CSV and JSON export of annotated review data, with optional
filtering by date range, model version, annotation status, and ground truth.
Can be called standalone or imported from other src/ modules.
"""
import csv
import io
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
def fetch_annotated(
db_connector,
annotated_only: bool = True,
model_version: Optional[str] = None,
date_from: Optional[str] = None,
date_to: Optional[str] = None,
ground_truth: Optional[bool] = None,
) -> List[Dict[str, Any]]:
"""Query the review_queue and return matching records as plain dicts."""
clauses: List[str] = []
params: List[Any] = []
if annotated_only:
clauses.append("rq.annotated = TRUE")
if model_version:
clauses.append("v.model_version = %s")
params.append(model_version)
if date_from:
clauses.append("rq.annotated_at >= %s")
params.append(date_from)
if date_to:
clauses.append("rq.annotated_at < %s")
params.append(date_to)
if ground_truth is not None:
clauses.append("rq.ground_truth = %s")
params.append(bool(ground_truth))
where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
rows = db_connector.fetchall(
f"""SELECT rq.video_id, v.file_path, rq.confidence_score, rq.routing_decision,
v.model_version, rq.ground_truth, rq.annotated_at, rq.notes,
pl.confidence_scores
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) rn
FROM processing_logs
) pl ON pl.video_id = rq.video_id AND pl.rn = 1
{where}
ORDER BY rq.annotated_at DESC""",
params if params else None,
)
return [_normalise(row) for row in (rows or [])]
def _normalise(row: dict) -> Dict[str, Any]:
try:
scores = json.loads(row["confidence_scores"]) if row.get("confidence_scores") else []
except (json.JSONDecodeError, TypeError):
scores = []
at = row.get("annotated_at")
return {
"video_id": row["video_id"],
"file_path": row["file_path"],
"confidence_score": row["confidence_score"],
"routing_decision": row["routing_decision"],
"model_version": row.get("model_version"),
"ground_truth": bool(row["ground_truth"]) if row["ground_truth"] is not None else None,
"annotated_at": at.isoformat() if isinstance(at, datetime) else at,
"notes": row.get("notes"),
"contributing_frames": scores,
}
def export_json(records: List[Dict[str, Any]], output_path: str) -> str:
"""Write records to a UTF-8 JSON file; returns the path."""
path = Path(output_path)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(records, default=str, indent=2), encoding="utf-8")
logger.info("Exported %d records to %s", len(records), path)
return str(path)
def export_csv(records: List[Dict[str, Any]], output_path: str) -> str:
"""Write records to a UTF-8 CSV file; returns the path."""
path = Path(output_path)
path.parent.mkdir(parents=True, exist_ok=True)
if not records:
path.write_text("", encoding="utf-8")
return str(path)
buf = io.StringIO()
writer = csv.DictWriter(buf, fieldnames=list(records[0].keys()))
writer.writeheader()
for r in records:
row = dict(r)
row["contributing_frames"] = json.dumps(row["contributing_frames"])
writer.writerow(row)
path.write_text(buf.getvalue(), encoding="utf-8")
logger.info("Exported %d records to %s", len(records), path)
return str(path)
-26
View File
@@ -1,26 +0,0 @@
"""Threshold-based routing of video-level confidence to MATCH / REVIEW / SKIP."""
import logging
logger = logging.getLogger(__name__)
MATCH = "MATCH"
REVIEW = "REVIEW"
SKIP = "SKIP"
def route(
confidence: float,
t_high: float = 0.75,
t_low: float = 0.45,
) -> str:
"""Return the routing decision for a video-level confidence score."""
if confidence >= t_high:
decision = MATCH
elif confidence >= t_low:
decision = REVIEW
else:
decision = SKIP
logger.debug("route: C=%.4f t_high=%.2f t_low=%.2f → %s", confidence, t_high, t_low, decision)
return decision
-537
View File
@@ -1,537 +0,0 @@
"""
Directory Scanner Service
Walks /data/input to discover new/modified files, computes hashes,
probes video metadata, validates codecs, and syncs state to MariaDB.
Supports incremental scanning for efficient 30TB corpus handling.
"""
import hashlib
import json
import logging
import os
import socket
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from codec_validator import CodecValidator
from prober import VideoMetadata, VideoProber
logger = logging.getLogger(__name__)
class DirectoryScanner:
"""Scan directories for new/modified video files and sync to DB."""
VIDEO_EXTENSIONS = {'.mp4', '.mov', '.avi', '.mkv', '.webm', '.flv', '.wmv'}
def __init__(
self,
db_connector,
config: dict,
input_path: str = "/data/input",
scan_interval: int = 60,
walker_threads: int = 8,
):
self.db = db_connector
self.config = config
self.input_path = Path(input_path)
self.scan_interval = scan_interval
self.walker_threads = walker_threads
self.prober = VideoProber(
timeout=config.get("scanner", {}).get("ffprobe_timeout_seconds", 10)
)
self.codec_validator = CodecValidator(
whitelist=set(config.get("codec", {}).get("whitelist", [])),
default_status_on_error=config.get("codec", {}).get("default_status_on_error", "UNSCANNABLE"),
)
self._running = False
self._scan_count = 0
self._total_files_discovered = 0
self._total_files_new = 0
self._total_files_modified = 0
self._total_files_unscannable = 0
# Single-instance guard: prevents overlapping scans both within this
# process (threading lock) and across replicas (DB lock with a lease).
self._instance_id = f"{socket.gethostname()}-{os.getpid()}-{uuid.uuid4().hex[:8]}"
self._lock_name = "directory_scanner"
self._lock_lease_seconds = int(
config.get("scanner", {}).get("lock_lease_seconds", 21600)
)
self._heartbeat_every = int(
config.get("scanner", {}).get("heartbeat_interval_files", 500)
)
self._scan_lock = threading.Lock()
self._lock_table_ensured = False
def start(self):
"""Start the scanner loop.
Each cycle is guarded so that at most one scan runs at a time:
- an in-process threading lock prevents re-entrant scans, and
- a DB lock (with a lease) prevents overlapping scans across replicas.
A long-running scan keeps its lease alive via heartbeats, so the next
scheduled tick (or another replica) waits instead of starting a second
parallel scan.
"""
self._running = True
logger.info(
"Scanner starting: input_path=%s interval=%ds threads=%d instance=%s",
self.input_path, self.scan_interval, self.walker_threads, self._instance_id,
)
while self._running:
# In-process re-entrancy guard: never run two scans at once.
if not self._scan_lock.acquire(blocking=False):
logger.warning("A scan is already in progress; skipping this cycle.")
self._sleep_interval()
continue
try:
if self.acquire_lock():
try:
self._run_scan()
finally:
self.release_lock()
else:
logger.info(
"Scanner lock held by another instance; skipping this cycle."
)
except Exception as e:
logger.error("Scanner error: %s", e, exc_info=True)
finally:
self._scan_lock.release()
# Sleep until next scan
self._sleep_interval()
def _sleep_interval(self):
"""Sleep for the scan interval, waking early if stopped."""
for _ in range(self.scan_interval):
if not self._running:
break
time.sleep(1)
def stop(self):
"""Stop the scanner."""
self._running = False
logger.info("Scanner stopping. Total scans: %d", self._scan_count)
# ------------------------------------------------------------------
# Single-instance lock (cross-process / cross-replica guard)
# ------------------------------------------------------------------
def _ensure_lock_table(self):
"""Create the scanner_lock table if it does not already exist."""
if self._lock_table_ensured:
return
self.db.execute(
"""CREATE TABLE IF NOT EXISTS scanner_lock (
lock_name VARCHAR(64) PRIMARY KEY,
owner VARCHAR(128) NOT NULL,
locked_at DATETIME NOT NULL,
lease_seconds INT NOT NULL DEFAULT 21600
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"""
)
self._lock_table_ensured = True
def acquire_lock(self) -> bool:
"""Attempt to acquire the cross-process scanner lock.
Returns True if this instance now owns the lock, False otherwise.
A stale lock (held longer than the lease) is taken over so a crashed
scanner does not block scanning forever.
"""
try:
self._ensure_lock_table()
except Exception as e:
logger.warning("Could not ensure scanner_lock table: %s", e)
return True # fail-open: keep scanning rather than block entirely
now = datetime.now(timezone.utc).replace(tzinfo=None)
# 1) Try to insert a fresh lock row.
try:
self.db.execute(
"""INSERT INTO scanner_lock (lock_name, owner, locked_at, lease_seconds)
VALUES (%s, %s, %s, %s)""",
(self._lock_name, self._instance_id, now, self._lock_lease_seconds),
transaction=True,
)
logger.info("Acquired scanner lock (fresh). owner=%s", self._instance_id)
return True
except Exception:
# Row already exists -> fall through to steal-if-stale.
pass
# 2) Take over the lock if it is stale or already ours.
stale_before = now - timedelta(seconds=self._lock_lease_seconds)
try:
affected = self.db.execute(
"""UPDATE scanner_lock
SET owner = %s, locked_at = %s
WHERE lock_name = %s
AND (owner = %s OR locked_at < %s)""",
(self._instance_id, now, self._lock_name, self._instance_id, stale_before),
transaction=True,
)
if affected and affected > 0:
logger.info("Acquired scanner lock (stale takeover). owner=%s",
self._instance_id)
return True
except Exception as e:
logger.warning("Failed to check scanner lock: %s", e)
return True # fail-open
logger.info("Scanner lock held by another instance; not acquiring.")
return False
def release_lock(self):
"""Release the scanner lock if we own it."""
try:
self.db.execute(
"""DELETE FROM scanner_lock WHERE lock_name = %s AND owner = %s""",
(self._lock_name, self._instance_id),
transaction=True,
)
except Exception as e:
logger.warning("Failed to release scanner lock: %s", e)
def _heartbeat(self):
"""Refresh the lock lease so a long-running scan is not stolen."""
try:
now = datetime.now(timezone.utc).replace(tzinfo=None)
self.db.execute(
"""UPDATE scanner_lock SET locked_at = %s
WHERE lock_name = %s AND owner = %s""",
(now, self._lock_name, self._instance_id),
transaction=True,
)
except Exception as e:
logger.debug("Scanner lock heartbeat failed: %s", e)
def _run_scan(self):
"""Execute a single scan cycle."""
scan_start = time.time()
self._scan_count += 1
logger.info("Scan #%d starting...", self._scan_count)
# Get last scan time from DB
last_scan = self._get_last_scan_time()
# Discover files
files_to_process, files_removed = self._discover_files(last_scan)
# Process files in parallel
results = self._process_files_batch(files_to_process)
# Update DB
scan_result = self._update_db(results, files_removed)
# Record scan history
duration = time.time() - scan_start
self._record_scan_history(scan_result, duration)
logger.info(
"Scan #%d complete: discovered=%d new=%d modified=%d unscannable=%d removed=%d duration=%.1fs",
self._scan_count,
scan_result["discovered"],
scan_result["new"],
scan_result["modified"],
scan_result["unscannable"],
len(files_removed),
duration,
)
def _get_last_scan_time(self) -> Optional[datetime]:
"""Get the last scan time from DB."""
result = self.db.fetchone(
"SELECT scan_end FROM scan_history ORDER BY id DESC LIMIT 1"
)
if result and result.get("scan_end"):
return result["scan_end"]
return None
def _discover_files(self, last_scan: Optional[datetime]) -> Tuple[List[Path], List[str]]:
"""Discover new and modified files in the input directory."""
files_to_process = []
files_removed = []
if not self.input_path.exists():
logger.warning("Input path does not exist: %s", self.input_path)
return files_to_process, files_removed
# Get files from DB for comparison
if last_scan:
# Only check files modified after last scan
db_files = self.db.fetchall(
"SELECT file_path, last_scan_time FROM videos WHERE last_scan_time > %s",
(last_scan,),
)
db_paths = {row["file_path"] for row in db_files}
# Check for modified files
for row in db_files:
file_path = Path(row["file_path"])
if file_path.exists() and file_path.is_file():
try:
mtime = datetime.fromtimestamp(
file_path.stat().st_mtime, tz=timezone.utc
)
if mtime > row["last_scan_time"]:
files_to_process.append(file_path)
except OSError:
pass
else:
# Full scan - walk the directory
logger.info("Full scan (no previous scan found). Walking %s...", self.input_path)
for root, dirs, filenames in os.walk(self.input_path):
for filename in filenames:
ext = Path(filename).suffix.lower()
if ext in self.VIDEO_EXTENSIONS:
file_path = Path(root) / filename
files_to_process.append(file_path)
# Check for removed files (if we have a last scan)
if last_scan:
all_db_files = self.db.fetchall(
"SELECT file_path FROM videos WHERE status NOT IN ('REMOVED', 'ERROR')",
)
current_paths = {str(f) for f in files_to_process}
for row in all_db_files:
fp = row["file_path"]
if fp not in current_paths and Path(fp).exists():
# File still exists but not in current walk - might be in a new directory
pass
elif fp not in current_paths and not Path(fp).exists():
files_removed.append(fp)
logger.info("Discovered %d files to process, %d files removed",
len(files_to_process), len(files_removed))
return files_to_process, files_removed
def _process_files_batch(self, files: List[Path]) -> List[dict]:
"""Process a batch of files in parallel."""
results = []
processed = 0
with ThreadPoolExecutor(max_workers=self.walker_threads) as executor:
future_to_file = {
executor.submit(self._process_single_file, f): f
for f in files
}
for future in as_completed(future_to_file):
file_path = future_to_file[future]
try:
result = future.result()
results.append(result)
except Exception as e:
logger.error("Error processing %s: %s", file_path, e, exc_info=True)
results.append({
"file_path": str(file_path),
"status": "ERROR",
"error_message": str(e),
})
processed += 1
# Keep the single-instance lock alive during long scans.
if self._heartbeat_every and processed % self._heartbeat_every == 0:
self._heartbeat()
return results
def _process_single_file(self, file_path: Path) -> dict:
"""Process a single file: hash, probe, validate, determine status."""
file_path_str = str(file_path)
result = {"file_path": file_path_str}
try:
# Compute hash
file_hash = self._compute_hash(file_path_str)
result["file_hash"] = file_hash
# Check for duplicate
existing = self.db.fetchone(
"SELECT id, status FROM videos WHERE file_hash = %s",
(file_hash,),
)
if existing:
result["status"] = existing["status"]
result["video_id"] = existing["id"]
result["action"] = "duplicate"
logger.debug("Duplicate file found: %s (id=%d, status=%s)",
file_path_str, existing["id"], existing["status"])
return result
# Probe video
metadata = self.prober.probe(file_path_str)
result.update(metadata.to_dict())
# Validate codec
if metadata.is_valid:
is_supported, reason = self.codec_validator.validate(metadata.codec)
if is_supported:
result["status"] = "PENDING"
result["action"] = "new"
self._total_files_new += 1
else:
result["status"] = "UNSCANNABLE"
result["error_message"] = f"Unsupported codec: {metadata.codec} ({reason})"
result["action"] = "unscannable"
self._total_files_unscannable += 1
logger.warning("Unsupported codec for %s: %s", file_path_str, metadata.codec)
else:
result["status"] = "UNSCANNABLE"
result["action"] = "unscannable"
self._total_files_unscannable += 1
except Exception as e:
result["status"] = "ERROR"
result["error_message"] = str(e)
logger.error("Error processing %s: %s", file_path_str, e, exc_info=True)
return result
def _compute_hash(self, file_path: str) -> str:
"""Compute SHA-256 hash of the first 1MB of a file."""
chunk_size = 1024 * 1024 # 1MB
sha256 = hashlib.sha256()
try:
with open(file_path, "rb") as f:
chunk = f.read(chunk_size)
if chunk:
sha256.update(chunk)
return sha256.hexdigest()
except (OSError, IOError) as e:
logger.error("Cannot hash file %s: %s", file_path, e)
return hashlib.sha256(file_path.encode()).hexdigest() # fallback
def _update_db(self, results: List[dict], files_removed: List[str]) -> dict:
"""Update database with scan results."""
stats = {"discovered": len(results), "new": 0, "modified": 0, "unscannable": 0}
# Batch insert new files
new_files = [r for r in results if r.get("action") == "new"]
if new_files:
self._batch_insert_new_files(new_files)
stats["new"] = len(new_files)
# Update unscannable files
unscannable = [r for r in results if r.get("action") == "unscannable" and r.get("video_id")]
if unscannable:
self._batch_update_status(unscannable, "UNSCANNABLE")
stats["unscannable"] = len(unscannable)
# Mark removed files
if files_removed:
self._mark_files_removed(files_removed)
stats["removed"] = len(files_removed)
return stats
def _batch_insert_new_files(self, files: List[dict]):
"""Batch insert new files into DB."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
for file_info in files:
try:
self.db.execute(
"""INSERT INTO videos
(file_path, file_hash, resolution_w, resolution_h, codec,
duration, status, last_scan_time, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""",
(
file_info["file_path"],
file_info["file_hash"],
file_info.get("resolution_w"),
file_info.get("resolution_h"),
file_info.get("codec"),
file_info.get("duration"),
"PENDING",
now,
now,
now,
),
transaction=True,
)
except Exception as e:
logger.error("Failed to insert %s: %s", file_info["file_path"], e)
def _batch_update_status(self, files: List[dict], status: str):
"""Batch update file status."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
for file_info in files:
video_id = file_info.get("video_id")
if video_id:
try:
self.db.execute(
"""UPDATE videos SET status = %s, error_message = %s,
last_scan_time = %s, updated_at = %s
WHERE id = %s""",
(status, file_info.get("error_message"), now, now, video_id),
transaction=True,
)
except Exception as e:
logger.error("Failed to update %s: %s", file_info["file_path"], e)
def _mark_files_removed(self, file_paths: List[str]):
"""Mark files as removed in DB."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
for fp in file_paths:
try:
self.db.execute(
"UPDATE videos SET status = 'REMOVED', last_scan_time = %s, updated_at = %s WHERE file_path = %s",
(now, now, fp),
transaction=True,
)
except Exception as e:
logger.error("Failed to mark %s as removed: %s", fp, e)
def _record_scan_history(self, stats: dict, duration: float):
"""Record scan history in DB."""
now = datetime.now(timezone.utc).replace(tzinfo=None)
self.db.execute(
"""INSERT INTO scan_history
(scan_start, scan_end, files_discovered, files_new, files_modified,
files_removed, files_unscannable, duration_seconds, status)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)""",
(
now,
now,
stats.get("discovered", 0),
stats.get("new", 0),
stats.get("modified", 0),
stats.get("removed", 0),
stats.get("unscannable", 0),
duration,
"COMPLETED",
),
transaction=True,
)
def get_stats(self) -> dict:
"""Get scanner statistics."""
return {
"scan_count": self._scan_count,
"total_files_discovered": self._total_files_discovered,
"total_files_new": self._total_files_new,
"total_files_modified": self._total_files_modified,
"total_files_unscannable": self._total_files_unscannable,
"instance_id": self._instance_id,
"scan_interval_seconds": self.scan_interval,
"lock_lease_seconds": self._lock_lease_seconds,
"scan_in_progress": self._scan_lock.locked(),
"input_path": str(self.input_path),
"is_running": self._running,
}
-42
View File
@@ -1,42 +0,0 @@
"""Utilities for managing scratch space used by frame sampling."""
import logging
import shutil
from pathlib import Path
from typing import Optional
logger = logging.getLogger(__name__)
class ScratchManager:
"""Create and clean per-video frame scratch directories."""
def __init__(self, base_path: str = "/scratch", video_id: Optional[str] = None, auto_cleanup: bool = True):
self.base_path = Path(base_path)
self.video_id = video_id or "unknown"
self.auto_cleanup = auto_cleanup
self.frame_dir = self.base_path / self.video_id / "frames"
def ensure_frame_dir(self) -> Path:
"""Create the per-video scratch directory if it does not exist."""
self.frame_dir.mkdir(parents=True, exist_ok=True)
return self.frame_dir
def cleanup(self) -> None:
"""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():
return 0
return sum(path.stat().st_size for path in self.frame_dir.rglob("*") if path.is_file())
-297
View File
@@ -1,297 +0,0 @@
"""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,
}
-186
View File
@@ -1,186 +0,0 @@
"""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())
-67
View File
@@ -1,67 +0,0 @@
import shutil
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from frame_sampler import FrameSampler, calculate_timestamps
from scratch_manager import ScratchManager
class Story03SamplingTests(unittest.TestCase):
def test_calculate_timestamps_uses_uniform_temporal_spacing(self):
stamps = calculate_timestamps(duration=90.0, interval=30.0)
self.assertEqual(stamps, [0.0, 30.0, 60.0])
def test_frame_sampler_builds_ffmpeg_command_with_scale_and_jpeg_output(self):
sampler = FrameSampler(interval_seconds=30, quality=2, output_format="jpeg")
command = sampler._build_ffmpeg_command(
video_path="/tmp/video.mp4",
output_path="/tmp/out.jpg",
timestamp=12.5,
resolution=(3840, 2160),
)
self.assertIn("ffmpeg", command[0])
self.assertIn("-ss", command)
self.assertIn("-vframes", command)
self.assertIn("scale=1920:1080", " ".join(command))
self.assertTrue(command[-1].endswith("out.jpg"))
def test_frame_sampler_uses_subprocess_and_returns_output_path(self):
sampler = FrameSampler(interval_seconds=30, quality=2, output_format="jpeg")
with tempfile.TemporaryDirectory() as tmpdir:
video_path = Path(tmpdir) / "sample.mp4"
output_path = Path(tmpdir) / "frame.jpg"
video_path.write_bytes(b"fake")
def _mock_run(*args, **kwargs):
output_path.write_bytes(b"frame")
return type("Completed", (), {"returncode": 0, "stdout": b"", "stderr": b""})()
with patch("subprocess.run", side_effect=_mock_run):
result = sampler.extract_frame(
video_path=str(video_path),
output_path=str(output_path),
timestamp=15.0,
resolution=(1280, 720),
)
self.assertTrue(result)
self.assertEqual(output_path.name, Path(result).name)
def test_scratch_manager_cleans_up_frames_after_processing(self):
with tempfile.TemporaryDirectory() as tmpdir:
manager = ScratchManager(base_path=tmpdir, video_id="video-1", auto_cleanup=True)
frame_dir = manager.ensure_frame_dir()
(frame_dir / "video-1_1000.jpg").write_bytes(b"frame")
manager.cleanup()
self.assertFalse(frame_dir.exists())
if __name__ == "__main__":
unittest.main()
-98
View File
@@ -1,98 +0,0 @@
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from batcher import DynamicBatcher
from face_detector import Detection, FaceDetector
from gpu_manager import GPUMemoryManager
class Story04FaceDetectionTests(unittest.TestCase):
def test_gpu_manager_respects_memory_thresholds(self):
manager = GPUMemoryManager(
max_memory_gb=18.0,
reduce_threshold_gb=16.0,
increase_threshold_gb=10.0,
initial_batch_size=16,
)
with patch.object(manager, "get_memory_stats", return_value=(17.0, 18.0)):
manager.adjust_batch_size()
self.assertLess(manager.current_batch_size, 16)
with patch.object(manager, "get_memory_stats", return_value=(8.0, 10.0)):
manager.adjust_batch_size()
self.assertGreater(manager.current_batch_size, manager.min_batch_size)
def test_dynamic_batcher_flushes_when_batch_full(self):
processed_batches = []
def process_fn(batch):
processed_batches.append(batch)
return [len(batch)]
batcher = DynamicBatcher(
process_fn=process_fn,
max_batch_size=3,
batch_timeout_ms=10000.0,
)
for i in range(5):
batcher.add(i)
self.assertEqual(len(processed_batches), 1)
self.assertEqual(processed_batches[0], [0, 1, 2])
self.assertEqual(batcher.queued_count, 2)
def test_face_detector_nms_removes_overlapping_boxes(self):
detector = FaceDetector(
engine_path="/nonexistent/model.trt",
confidence_threshold=0.1,
iou_threshold=0.45,
max_faces_per_frame=10,
)
duplicates = [
Detection("frame.jpg", 10, 10, 50, 50, 0.9),
Detection("frame.jpg", 12, 12, 48, 48, 0.8),
Detection("frame.jpg", 100, 100, 150, 150, 0.75),
]
kept = detector._nms(duplicates)
self.assertEqual(len(kept), 2)
self.assertAlmostEqual(kept[0].confidence, 0.9, places=5)
def test_face_detector_extracts_and_resizes_crops(self):
detector = FaceDetector(
engine_path="/nonexistent/model.trt",
input_size=640,
confidence_threshold=0.25,
)
with tempfile.TemporaryDirectory() as tmpdir:
from PIL import Image
frame_path = Path(tmpdir) / "frame.jpg"
image = Image.new("RGB", (640, 480), color=(100, 150, 200))
image.save(frame_path)
detections = [
Detection(str(frame_path), 0, 0, 640, 480, 0.95),
]
cropped = detector.extract_crops(detections, output_dir=tmpdir, crop_size=(224, 224))
self.assertEqual(len(cropped), 1)
self.assertTrue(Path(cropped[0].crop_path).exists())
with Image.open(cropped[0].crop_path) as crop:
self.assertEqual(crop.size, (224, 224))
if __name__ == "__main__":
unittest.main()
-83
View File
@@ -1,83 +0,0 @@
import sys
import unittest
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import aggregator
import router
from classifier import FaceClassifier, calibrated_softmax
class Story05ClassifierTests(unittest.TestCase):
def test_calibrated_softmax_sums_to_one(self):
logits = np.array([[2.0, 1.0], [-1.0, 3.0]], dtype=np.float32)
probs = calibrated_softmax(logits, temperature=1.0)
np.testing.assert_allclose(probs.sum(axis=1), [1.0, 1.0], atol=1e-6)
def test_temperature_scaling_raises_lower_confidence_entropy(self):
logits = np.array([[2.0, 0.5]], dtype=np.float32)
sharp = calibrated_softmax(logits, temperature=0.5)
soft = calibrated_softmax(logits, temperature=2.0)
# higher temperature → softer distribution (target class prob moves toward 0.5)
self.assertGreater(sharp[0, 0], soft[0, 0])
def test_classifier_placeholder_returns_neutral_probability(self):
clf = FaceClassifier(engine_path="/nonexistent/model.trt", temperature=1.0)
probs = clf.classify([])
self.assertEqual(probs, [])
def test_classifier_placeholder_single_crop_returns_half(self):
import tempfile
from PIL import Image
clf = FaceClassifier(engine_path="/nonexistent/model.trt", temperature=1.0)
with tempfile.TemporaryDirectory() as tmpdir:
crop = Path(tmpdir) / "crop.jpg"
Image.new("RGB", (224, 224)).save(crop)
probs = clf.classify([str(crop)])
# placeholder logits are all zeros → softmax → 0.5 for each class
self.assertAlmostEqual(probs[0], 0.5, places=5)
class Story05AggregatorTests(unittest.TestCase):
def test_max_strategy(self):
self.assertAlmostEqual(aggregator.aggregate([0.3, 0.8, 0.6], strategy="max"), 0.8)
def test_empty_confidences_returns_zero(self):
self.assertEqual(aggregator.aggregate([], strategy="max"), 0.0)
def test_top_k_mean(self):
result = aggregator.aggregate([0.1, 0.9, 0.5, 0.8], strategy="top_k_mean", top_k=2)
self.assertAlmostEqual(result, (0.9 + 0.8) / 2, places=5)
def test_weighted_mean_clamps_to_unit_interval(self):
result = aggregator.aggregate([1.0, 1.0], strategy="weighted_mean", alpha=100.0, beta=0.0)
self.assertLessEqual(result, 1.0)
self.assertGreaterEqual(result, 0.0)
class Story05RouterTests(unittest.TestCase):
def test_match_at_high_threshold(self):
self.assertEqual(router.route(0.75), router.MATCH)
def test_review_between_thresholds(self):
self.assertEqual(router.route(0.60), router.REVIEW)
def test_skip_below_low_threshold(self):
self.assertEqual(router.route(0.44), router.SKIP)
def test_inclusive_high_threshold_boundary(self):
self.assertEqual(router.route(0.75, t_high=0.75, t_low=0.45), router.MATCH)
def test_inclusive_low_threshold_boundary(self):
self.assertEqual(router.route(0.45, t_high=0.75, t_low=0.45), router.REVIEW)
def test_no_faces_zero_confidence_routes_skip(self):
self.assertEqual(router.route(0.0), router.SKIP)
if __name__ == "__main__":
unittest.main()
-148
View File
@@ -1,148 +0,0 @@
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()
-95
View File
@@ -1,95 +0,0 @@
import csv
import io
import json
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import review_export
def _make_db(rows):
db = MagicMock()
db.fetchall.return_value = rows
return db
class Story07ReviewExportTests(unittest.TestCase):
_ROWS = [
{
"video_id": 1, "file_path": "/data/input/a.mp4",
"confidence_score": 0.62, "routing_decision": "REVIEW",
"model_version": "v1.0", "ground_truth": True,
"annotated_at": None, "notes": "ok",
"confidence_scores": json.dumps([0.60, 0.62, 0.65]),
},
]
def test_fetch_annotated_passes_correct_where_clause(self):
db = _make_db(self._ROWS)
review_export.fetch_annotated(db, annotated_only=True, model_version="v1.0")
call_args = db.fetchall.call_args
sql = call_args[0][0]
self.assertIn("rq.annotated = TRUE", sql)
self.assertIn("v.model_version = %s", sql)
def test_fetch_annotated_normalises_confidence_scores(self):
db = _make_db(self._ROWS)
records = review_export.fetch_annotated(db)
self.assertIsInstance(records[0]["contributing_frames"], list)
self.assertEqual(records[0]["contributing_frames"], [0.60, 0.62, 0.65])
def test_fetch_annotated_handles_null_confidence_scores(self):
rows = [{**self._ROWS[0], "confidence_scores": None}]
db = _make_db(rows)
records = review_export.fetch_annotated(db)
self.assertEqual(records[0]["contributing_frames"], [])
def test_export_json_writes_valid_utf8_file(self):
with tempfile.TemporaryDirectory() as tmpdir:
records = review_export.fetch_annotated(_make_db(self._ROWS))
path = review_export.export_json(records, f"{tmpdir}/out.json")
loaded = json.loads(Path(path).read_text(encoding="utf-8"))
self.assertEqual(len(loaded), 1)
self.assertEqual(loaded[0]["video_id"], 1)
self.assertTrue(loaded[0]["ground_truth"])
def test_export_csv_writes_valid_csv(self):
with tempfile.TemporaryDirectory() as tmpdir:
records = review_export.fetch_annotated(_make_db(self._ROWS))
path = review_export.export_csv(records, f"{tmpdir}/out.csv")
content = Path(path).read_text(encoding="utf-8")
reader = csv.DictReader(io.StringIO(content))
rows = list(reader)
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]["routing_decision"], "REVIEW")
# contributing_frames should be a JSON string in CSV
frames = json.loads(rows[0]["contributing_frames"])
self.assertEqual(frames, [0.60, 0.62, 0.65])
def test_export_csv_empty_returns_empty_file(self):
with tempfile.TemporaryDirectory() as tmpdir:
path = review_export.export_csv([], f"{tmpdir}/empty.csv")
self.assertEqual(Path(path).read_text(), "")
def test_ground_truth_filter_appears_in_query(self):
db = _make_db([])
review_export.fetch_annotated(db, ground_truth=False)
sql = db.fetchall.call_args[0][0]
self.assertIn("rq.ground_truth = %s", sql)
class Story07AppSyntaxTest(unittest.TestCase):
def test_app_module_compiles(self):
"""Ensure ui/app.py has no syntax errors."""
app_path = Path(__file__).resolve().parents[1] / "ui" / "app.py"
source = app_path.read_text(encoding="utf-8")
compile(source, str(app_path), "exec")
if __name__ == "__main__":
unittest.main()
-550
View File
@@ -1,550 +0,0 @@
"""
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()
-521
View File
@@ -1,521 +0,0 @@
"""
Tests for Story 09: Observability, Monitoring & Hardening.
Covers:
- metrics.py: NoOp fallback, update_queue_depths, update_scratch_metrics
- crash_recovery.py: recover_on_startup, checkpointing, idempotency guard
- retry.py: successful call, retry on transient error, non-retryable bypass, exhaustion
- drift_detector.py: detect_drift logic, all alert checks
- health_check.py: HealthStatus snapshot, HTTP /health endpoint
"""
import json
import os
import sys
import tempfile
import threading
import time
import unittest
import urllib.request
from pathlib import Path
from unittest.mock import MagicMock, call, patch
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from crash_recovery import CrashRecovery
from drift_detector import DriftDetector, detect_drift
from health_check import HealthStatus, start_health_server, health
from retry import RetryExhaustedError, retry, _is_non_retryable
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_config(overrides: dict = None):
cfg = MagicMock()
mon_defaults = {
"alerts": {
"confidence_drift_threshold": 0.10,
"review_queue_max_size": 1000,
"review_queue_max_age_hours": 24,
"throughput_min_videos_per_hour": 20,
"throughput_min_duration_hours": 1,
"error_rate_threshold": 0.05,
"error_rate_window_hours": 1,
},
"drift_detection": {"enabled": True, "schedule": "0 2 * * 0", "baseline_source": "db"},
"crash_recovery": {"lock_timeout_minutes": 5, "auto_requeue": True},
"retry": {"max_attempts": 3, "initial_delay": 0.0, "backoff_factor": 2.0},
}
if overrides:
mon_defaults.update(overrides)
def _get_section(section):
if section == "monitoring":
return mon_defaults
return {}
def _get(path, default=None):
mapping = {
"storage.scratch_path": "/tmp/test_scratch",
"storage.models_path": "/tmp/test_models",
}
return mapping.get(path, default)
cfg.get_section.side_effect = _get_section
cfg.get.side_effect = _get
return cfg
def _make_db():
return MagicMock()
# ---------------------------------------------------------------------------
# metrics.py tests
# ---------------------------------------------------------------------------
class TestMetricsNoOp(unittest.TestCase):
"""The _NoOpMetric must absorb all method calls without raising."""
def test_noop_labels_inc_does_not_raise(self):
from metrics import _NoOpMetric
m = _NoOpMetric()
m.labels(routing_decision="MATCH").inc()
def test_noop_set_does_not_raise(self):
from metrics import _NoOpMetric
m = _NoOpMetric()
m.set(42)
def test_noop_observe_does_not_raise(self):
from metrics import _NoOpMetric
m = _NoOpMetric()
m.observe(0.75)
class TestMetricsQueueDepths(unittest.TestCase):
def test_update_queue_depths_sets_gauges(self):
from metrics import update_queue_depths, queue_depth_pending, queue_depth_processing
db = MagicMock()
db.fetchall.side_effect = [
[{"status": "PENDING", "cnt": 10}, {"status": "PROCESSING", "cnt": 3}],
[{"cnt": 7}],
]
# Should not raise even if prometheus is absent
update_queue_depths(db)
def test_update_queue_depths_handles_db_error_gracefully(self):
from metrics import update_queue_depths
db = MagicMock()
db.fetchall.side_effect = Exception("DB down")
update_queue_depths(db) # must not raise
class TestMetricsScratch(unittest.TestCase):
def test_update_scratch_metrics_runs_without_error(self):
from metrics import update_scratch_metrics
with tempfile.TemporaryDirectory() as tmpdir:
update_scratch_metrics(tmpdir) # must not raise
def test_update_scratch_metrics_handles_missing_path(self):
from metrics import update_scratch_metrics
update_scratch_metrics("/nonexistent_path_xyz") # must not raise
# ---------------------------------------------------------------------------
# crash_recovery.py tests
# ---------------------------------------------------------------------------
class TestCrashRecoveryRequeue(unittest.TestCase):
def test_recover_on_startup_calls_update(self):
db = _make_db()
db.execute.return_value = 3
cfg = _make_config()
cr = CrashRecovery(db, cfg)
count = cr.recover_on_startup()
self.assertEqual(count, 3)
db.execute.assert_called_once()
sql = db.execute.call_args[0][0]
self.assertIn("PENDING", sql)
self.assertIn("PROCESSING", sql)
def test_recover_on_startup_skipped_when_disabled(self):
db = _make_db()
cfg = _make_config({"crash_recovery": {"lock_timeout_minutes": 5, "auto_requeue": False}})
cr = CrashRecovery(db, cfg)
count = cr.recover_on_startup()
self.assertEqual(count, 0)
db.execute.assert_not_called()
def test_recover_on_startup_handles_db_error(self):
db = _make_db()
db.execute.side_effect = Exception("connection refused")
cfg = _make_config()
cr = CrashRecovery(db, cfg)
count = cr.recover_on_startup()
self.assertEqual(count, 0)
def test_list_stuck_videos_returns_rows(self):
db = _make_db()
db.fetchall.return_value = [{"id": 5, "file_path": "/data/vid.mp4", "updated_at": None}]
cfg = _make_config()
cr = CrashRecovery(db, cfg)
rows = cr.list_stuck_videos()
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]["id"], 5)
class TestCrashRecoveryCheckpoint(unittest.TestCase):
def test_save_and_load_checkpoint_roundtrip(self):
with tempfile.TemporaryDirectory() as tmpdir:
cfg = _make_config()
cfg.get.side_effect = lambda k, d=None: {
"storage.scratch_path": tmpdir,
}.get(k, d)
cr = CrashRecovery(_make_db(), cfg)
cr.save_checkpoint(42, "extracting", {"frames_done": 5})
result = cr.load_checkpoint(42)
self.assertIsNotNone(result)
self.assertEqual(result["video_id"], 42)
self.assertEqual(result["state"], "extracting")
self.assertEqual(result["progress"]["frames_done"], 5)
def test_load_checkpoint_returns_none_when_absent(self):
with tempfile.TemporaryDirectory() as tmpdir:
cfg = _make_config()
cfg.get.side_effect = lambda k, d=None: {
"storage.scratch_path": tmpdir,
}.get(k, d)
cr = CrashRecovery(_make_db(), cfg)
self.assertIsNone(cr.load_checkpoint(999))
def test_delete_checkpoint_removes_file(self):
with tempfile.TemporaryDirectory() as tmpdir:
cfg = _make_config()
cfg.get.side_effect = lambda k, d=None: {
"storage.scratch_path": tmpdir,
}.get(k, d)
cr = CrashRecovery(_make_db(), cfg)
cr.save_checkpoint(7, "classifying", {})
cr.delete_checkpoint(7)
self.assertIsNone(cr.load_checkpoint(7))
def test_checkpoint_timestamp_is_iso_format(self):
with tempfile.TemporaryDirectory() as tmpdir:
cfg = _make_config()
cfg.get.side_effect = lambda k, d=None: {
"storage.scratch_path": tmpdir,
}.get(k, d)
cr = CrashRecovery(_make_db(), cfg)
cr.save_checkpoint(1, "detecting", {})
ckpt = cr.load_checkpoint(1)
# Should parse without error
from datetime import datetime
datetime.fromisoformat(ckpt["timestamp"].replace("Z", "+00:00"))
class TestIdempotencyGuard(unittest.TestCase):
def test_returns_true_when_already_completed(self):
db = _make_db()
db.fetchall.return_value = [{"status": "COMPLETED"}]
cr = CrashRecovery(db, _make_config())
self.assertTrue(cr.is_already_completed(1))
def test_returns_false_when_not_completed(self):
db = _make_db()
db.fetchall.return_value = [{"status": "PENDING"}]
cr = CrashRecovery(db, _make_config())
self.assertFalse(cr.is_already_completed(1))
def test_returns_false_when_no_row(self):
db = _make_db()
db.fetchall.return_value = []
cr = CrashRecovery(db, _make_config())
self.assertFalse(cr.is_already_completed(99))
# ---------------------------------------------------------------------------
# retry.py tests
# ---------------------------------------------------------------------------
class TestRetryDecorator(unittest.TestCase):
def test_successful_call_returns_value(self):
@retry(max_attempts=3, initial_delay=0.0, step="test")
def always_succeeds():
return 42
self.assertEqual(always_succeeds(), 42)
def test_retries_on_transient_error(self):
call_count = {"n": 0}
@retry(max_attempts=3, initial_delay=0.0, step="test")
def flaky():
call_count["n"] += 1
if call_count["n"] < 3:
raise ConnectionError("transient")
return "ok"
result = flaky()
self.assertEqual(result, "ok")
self.assertEqual(call_count["n"], 3)
def test_raises_retry_exhausted_after_max_attempts(self):
@retry(max_attempts=3, initial_delay=0.0, step="test")
def always_fails():
raise ConnectionError("always fails")
with self.assertRaises(RetryExhaustedError):
always_fails()
def test_non_retryable_error_propagates_immediately(self):
call_count = {"n": 0}
@retry(max_attempts=3, initial_delay=0.0, step="test")
def raises_non_retryable():
call_count["n"] += 1
raise FileNotFoundError("no such file")
with self.assertRaises(FileNotFoundError):
raises_non_retryable()
self.assertEqual(call_count["n"], 1)
def test_only_specified_exception_types_are_retried(self):
@retry(max_attempts=3, initial_delay=0.0, exceptions=(ValueError,), step="test")
def raises_type_error():
raise TypeError("wrong type")
with self.assertRaises(TypeError):
raises_type_error()
def test_preserves_return_value_on_first_try(self):
@retry(max_attempts=5, initial_delay=0.0, step="test")
def returns_dict():
return {"key": "value"}
self.assertEqual(returns_dict(), {"key": "value"})
class TestIsNonRetryable(unittest.TestCase):
def test_file_not_found_is_non_retryable(self):
self.assertTrue(_is_non_retryable(FileNotFoundError("x")))
def test_permission_error_is_non_retryable(self):
self.assertTrue(_is_non_retryable(PermissionError("x")))
def test_connection_error_is_retryable(self):
self.assertFalse(_is_non_retryable(ConnectionError("x")))
def test_runtime_error_is_retryable(self):
self.assertFalse(_is_non_retryable(RuntimeError("x")))
# ---------------------------------------------------------------------------
# drift_detector.py tests
# ---------------------------------------------------------------------------
class TestDetectDrift(unittest.TestCase):
def test_no_drift_when_distributions_match(self):
base = [0.3] * 50 + [0.7] * 50 # 50% high confidence
curr = [0.3] * 50 + [0.7] * 50
self.assertFalse(detect_drift(curr, base, threshold=0.10))
def test_drift_detected_when_shift_exceeds_threshold(self):
base = [0.3] * 80 + [0.8] * 20 # 20% high
curr = [0.8] * 70 + [0.3] * 30 # 70% high → shift = 0.50
self.assertTrue(detect_drift(curr, base, threshold=0.10))
def test_no_drift_just_below_threshold(self):
base = [0.8] * 50 + [0.2] * 50 # 50% high
curr = [0.8] * 59 + [0.2] * 41 # 59% high → shift = 9%
self.assertFalse(detect_drift(curr, base, threshold=0.10))
def test_drift_at_boundary(self):
base = [0.8] * 50 + [0.2] * 50 # 50%
curr = [0.8] * 61 + [0.2] * 39 # 61% → shift = 11%
self.assertTrue(detect_drift(curr, base, threshold=0.10))
def test_empty_current_returns_false(self):
self.assertFalse(detect_drift([], [0.5] * 10, threshold=0.10))
def test_empty_baseline_returns_false(self):
self.assertFalse(detect_drift([0.5] * 10, [], threshold=0.10))
class TestDriftDetectorAlerts(unittest.TestCase):
def _make_detector(self, db=None, overrides=None):
alerts = []
cfg = _make_config(overrides or {})
db = db or _make_db()
detector = DriftDetector(db, cfg, alert_fn=alerts.append)
return detector, alerts
def test_check_review_queue_growth_triggers_alert(self):
db = _make_db()
db.fetchall.return_value = [{"cnt": 1500}]
detector, alerts = self._make_detector(db)
triggered = detector.check_review_queue_growth()
self.assertTrue(triggered)
self.assertEqual(len(alerts), 1)
self.assertIn("1500", alerts[0])
def test_check_review_queue_growth_no_alert_below_threshold(self):
db = _make_db()
db.fetchall.return_value = [{"cnt": 50}]
detector, alerts = self._make_detector(db)
triggered = detector.check_review_queue_growth()
self.assertFalse(triggered)
self.assertEqual(len(alerts), 0)
def test_check_low_throughput_triggers_alert(self):
db = _make_db()
db.fetchall.return_value = [{"cnt": 5}] # 5 videos in last 1h < 20 min
detector, alerts = self._make_detector(db)
triggered = detector.check_low_throughput()
self.assertTrue(triggered)
self.assertEqual(len(alerts), 1)
def test_check_low_throughput_no_alert_above_threshold(self):
db = _make_db()
db.fetchall.return_value = [{"cnt": 50}] # 50 > 20
detector, alerts = self._make_detector(db)
triggered = detector.check_low_throughput()
self.assertFalse(triggered)
def test_check_error_rate_triggers_alert(self):
db = _make_db()
db.fetchall.return_value = [
{"status": "COMPLETED", "cnt": 80},
{"status": "ERROR", "cnt": 10},
{"status": "UNSCANNABLE", "cnt": 10},
]
detector, alerts = self._make_detector(db)
triggered = detector.check_error_rate()
self.assertTrue(triggered) # 20/100 = 20% > 5%
self.assertEqual(len(alerts), 1)
def test_check_error_rate_no_alert_below_threshold(self):
db = _make_db()
db.fetchall.return_value = [
{"status": "COMPLETED", "cnt": 98},
{"status": "ERROR", "cnt": 2},
]
detector, alerts = self._make_detector(db)
triggered = detector.check_error_rate()
self.assertFalse(triggered) # 2% < 5%
def test_check_error_rate_no_alert_zero_videos(self):
db = _make_db()
db.fetchall.return_value = []
detector, alerts = self._make_detector(db)
triggered = detector.check_error_rate()
self.assertFalse(triggered)
def test_run_all_checks_returns_dict_with_expected_keys(self):
db = _make_db()
db.fetchall.return_value = [{"cnt": 0}]
detector, _ = self._make_detector(db)
with patch.object(detector, "_fetch_recent_confidences", return_value=[]):
results = detector.run_all_checks()
self.assertIn("confidence_drift", results)
self.assertIn("review_queue_growth", results)
self.assertIn("low_throughput", results)
self.assertIn("high_error_rate", results)
def test_check_handles_db_error_gracefully(self):
db = _make_db()
db.fetchall.side_effect = Exception("DB offline")
detector, alerts = self._make_detector(db)
# Should not raise
self.assertFalse(detector.check_review_queue_growth())
self.assertFalse(detector.check_low_throughput())
self.assertFalse(detector.check_error_rate())
# ---------------------------------------------------------------------------
# health_check.py tests
# ---------------------------------------------------------------------------
class TestHealthStatus(unittest.TestCase):
def test_snapshot_contains_required_keys(self):
hs = HealthStatus()
snap = hs.snapshot()
for key in ("status", "gpu_available", "gpu_memory_used_gb",
"queue_depth", "uptime_seconds", "videos_processed_today", "last_error"):
self.assertIn(key, snap, f"Missing key: {key}")
def test_update_changes_values(self):
hs = HealthStatus()
hs.update(status="healthy", queue_depth=55)
snap = hs.snapshot()
self.assertEqual(snap["status"], "healthy")
self.assertEqual(snap["queue_depth"], 55)
def test_uptime_increases_over_time(self):
hs = HealthStatus()
snap1 = hs.snapshot()
time.sleep(0.05)
snap2 = hs.snapshot()
self.assertGreaterEqual(snap2["uptime_seconds"], snap1["uptime_seconds"])
def test_update_is_thread_safe(self):
hs = HealthStatus()
errors = []
def writer(n):
try:
for _ in range(100):
hs.update(queue_depth=n)
except Exception as exc:
errors.append(exc)
threads = [threading.Thread(target=writer, args=(i,)) for i in range(5)]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(errors, [])
def test_http_health_endpoint_returns_200(self):
"""Start a real health server and hit /health with urllib."""
import socket
# Find a free port
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
t = start_health_server(port)
time.sleep(0.1) # give the server a moment to bind
url = f"http://127.0.0.1:{port}/health"
with urllib.request.urlopen(url, timeout=2) as resp:
self.assertEqual(resp.status, 200)
body = json.loads(resp.read())
self.assertIn("status", body)
self.assertIn("uptime_seconds", body)
def test_http_404_for_unknown_path(self):
"""Non /health paths return 404."""
import socket
from urllib.error import HTTPError
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
start_health_server(port)
time.sleep(0.1)
with self.assertRaises(HTTPError) as ctx:
urllib.request.urlopen(f"http://127.0.0.1:{port}/unknown", timeout=2)
self.assertEqual(ctx.exception.code, 404)
if __name__ == "__main__":
unittest.main()
+407 -469
View File
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -1,4 +1,4 @@
Flask==3.0.0 Flask==3.0.0
gunicorn==21.2.0 gunicorn==21.2.0
PyMySQL==1.1.0
Jinja2==3.1.2 Jinja2==3.1.2
requests>=2.28.0
-68
View File
@@ -1,68 +0,0 @@
# Base image: CUDA 11.8 runtime on Ubuntu 22.04
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
# Avoid interactive prompts during build
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC
# System dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
# Python
python3.10 \
python3.10-dev \
python3.10-venv \
python3-pip \
# FFmpeg
ffmpeg \
# OpenCV dependencies
libgl1-mesa-glx \
libglib2.0-0 \
libsm6 \
libxext6 \
libxrender-dev \
# Utilities
curl \
wget \
git \
ca-certificates \
&& rm -rf /var/lib/apt/lists/*
# Create non-root user
RUN groupadd -g 1000 appuser && \
useradd -m -u 1000 -g appuser -s /bin/bash appuser
# Set working directory
WORKDIR /app
# Create necessary directories
RUN mkdir -p /scratch /models /data/training /data/output /logs && \
chown -R appuser:appuser /app /scratch /models /data /logs
# Copy requirements first for better caching
COPY worker/requirements.txt /app/requirements.txt
# Create and activate virtual environment
RUN python3.10 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
# Install Python dependencies
RUN pip install --no-cache-dir -r requirements.txt
# Copy application code
COPY src/ /app/src/
COPY config.yaml /app/config.yaml
# Set environment variables
ENV PYTHONUNBUFFERED=1
ENV PATH="/opt/venv/bin:$PATH"
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
# Health check
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
CMD python3 -c "import torch; assert torch.cuda.is_available()" || exit 1
# Switch to non-root user
USER appuser
# Default command: process tasks from the API until drained
CMD ["python3", "-m", "src.task_worker"]
-40
View File
@@ -1,40 +0,0 @@
# 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
-40
View File
@@ -1,40 +0,0 @@
# Core ML Frameworks
torch==2.1.0
--extra-index-url https://download.pytorch.org/whl/cu118
# TensorRT (compatible with CUDA 11.8)
# Note: Install from NVIDIA repo or wheel for your specific platform
# tensorrt==8.6.1 # Uncomment after verifying compatibility
# ONNX Runtime with GPU support
onnxruntime-gpu==1.16.3
# Computer Vision
opencv-python-headless==4.8.1.78
Pillow==10.1.0
# Database
PyMySQL==1.1.0
DBUtils==3.1.0
# Data Processing
numpy==1.24.4
pandas==2.1.4
pyarrow==14.0.0
# Configuration
pyyaml==6.0.1
# Logging
python-json-logger==2.0.7
# Monitoring
prometheus-client==0.19.0
# Utilities
tqdm==4.66.1
filelock==3.13.1
# Testing (dev dependencies - install with [dev] extra)
# pytest==7.4.3
# pytest-cov==4.1.0