thebears/pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

YACWC

Browse Source
This commit is contained in:
Ishan S. Patel
2026-07-08 10:53:37 -04:00
parent 19ace202ba
commit d286b9b311
25 changed files with 527322 additions and 443 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vision_v3/30_queued_upload_to_db/30_queued_upload_to_ch.py
+202
View File
@@ -0,0 +1,202 @@
import numpy as np
import time
from multiprocessing import Process, Queue, Manager
import threading
from common_code import kwq, util, file_names, vector_utils, video_meta
from datetime import datetime
from clickhouse_connect import get_client
topic_subscribe = 'upload_to_ch_db'
group_id = "vision_v3_uploaders"
def get_score_file(f):
return f.key().decode().strip('"')
def insert_into_ch(file_path, client):
cset = util.get_cset_match(file_path)
ff = np.load(cset['.detection.npz'], allow_pickle = True)
det_results = dict(ff)
del_idx = list()
filt_scores = det_results['final_scores']
filt_frames = det_results['final_frames']
spec_label_idx = ff['species_label_map'].tolist()
ftp_vid = video_meta.FTPVideo(cset)
frame_num_off = ftp_vid.sorted_frames_info
frame_period = ftp_vid.frameperiod
frame_rate = ftp_vid.framerate
indexes = np.searchsorted(ff['scored_frames'], filt_frames)
dff = np.diff(ff['scored_frames'])
head = np.insert(dff, 0, 0)
tail = np.append(dff, 0)
all_frame_durations = frame_period*np.mean(np.stack([head,tail]), axis=0)
frame_file = [ftp_vid.cpath.split('/')[-1].replace('_reduced','').replace('.mp4','')]*len(filt_frames)
frame_bboxes = ff['final_boxes']
frame_numbers = filt_frames
frame_scores = filt_scores
frame_species = [spec_label_idx[int(x)] for x in ff['final_labels']]
frame_durations = all_frame_durations[indexes]
frame_tstamps = [frame_num_off[x]['time'] for x in filt_frames]
db_insert_times = [datetime.now()]*len(filt_frames)
a = frame_numbers
starts = np.r_[0, np.flatnonzero(a[1:] != a[:-1]) + 1]
det_count_num = np.arange(len(a)) - np.repeat(starts, np.diff(np.r_[starts, len(a)]))
camera = [ftp_vid.camera_name] * len(filt_frames)
result = client.insert('nuggets.cameras_object_detection_results',
column_oriented = True,
column_names = ['camera','video_filename','frame_number','detection_number','classification_score','object_name','bbox_xyxy','frame_duration','frame_timestamp','db_insert_time'],
data = [camera, frame_file, frame_numbers, det_count_num, frame_scores, frame_species, frame_bboxes.tolist(), frame_durations, frame_tstamps, db_insert_times]
)
# %%
def watch_and_upload(rank, health_queue):
client = get_client(host='192.168.1.242', compression=True)
start_time = time.time()
client_suffix = '_'+str(rank)
client_id = 'uploader_to_ch_db'+client_suffix
current_time = time.time()
health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
queue = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'earliest')
queue.subscribe([topic_subscribe])
last_heartbeat = time.time()
while True:
current_time = time.time()
# # Exit after 5 minutes
# if (current_time - start_time) > 60*5:
# print('###########################')
# print('######### EXITING #########')
# print('###########################')
# return
# Send heartbeat every 10 seconds
if current_time - last_heartbeat > 5:
try:
health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
last_heartbeat = current_time
except:
# Queue might be full, continue working
pass
result = queue.poll(timeout=0.05)
if result is not None:
try:
score_this = get_score_file(result)
insert_into_ch(score_this, client)
print('#' + str(rank) + ' Uploaded to ' + score_this)
# Send success heartbeat
try:
health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'working'})
except:
pass
except Exception as e:
print('Failed ' + score_this + ':' + str(e))
queue.commit()
def health_monitor(processes, health_queues, num_workers):
"""Monitor process health and restart hung processes"""
last_heartbeat = {}
HEALTH_TIMEOUT = 120 # Consider process hung if no heartbeat for 60 seconds
while True:
time.sleep(5) # Check every 5 seconds
current_time = time.time()
for rank in range(num_workers):
# Drain the health queue for this process
latest_heartbeat = None
try:
while True:
try:
heartbeat = health_queues[rank].get_nowait()
if heartbeat['rank'] == rank:
latest_heartbeat = heartbeat['timestamp']
except:
break
except:
pass
# Update last known heartbeat
if latest_heartbeat:
last_heartbeat[rank] = latest_heartbeat
# Check if process is alive and responsive
process = processes[rank]
last_seen = last_heartbeat.get(rank, 0)
if not process.is_alive():
print(f"Process {rank} has died. Restarting...")
processes[rank] = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
processes[rank].start()
last_heartbeat[rank] = current_time
elif current_time - last_seen > HEALTH_TIMEOUT:
print(f"Process {rank} appears hung (no heartbeat for {current_time - last_seen:.1f}s). Restarting...")
process.terminate()
process.join(timeout=5)
if process.is_alive():
process.kill()
process.join()
# Create new process
processes[rank] = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
processes[rank].start()
last_heartbeat[rank] = current_time
# Main execution
if __name__ == "__main__":
num_workers = 4
# Create health queues for each worker
health_queues = [Queue(maxsize=10) for _ in range(num_workers)]
# Create and start worker processes
processes = []
for rank in range(num_workers):
p = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
p.start()
processes.append(p)
# Start health monitor in a separate thread
monitor_thread = threading.Thread(
target=health_monitor,
args=(processes, health_queues, num_workers),
daemon=True
)
monitor_thread.start()
try:
# Wait for all processes to complete (they will exit after 5 minutes)
for p in processes:
p.join()
print("All worker processes completed")
except KeyboardInterrupt:
print("Shutting down...")
for p in processes:
p.terminate()
p.join(timeout=5)
if p.is_alive():
p.kill()
p.join()
vision_v3/30_queued_upload_to_db/30_queued_upload_to_db.py
+142
View File
@@ -0,0 +1,142 @@
from common_code.vector_utils import upload_vectors_to_db
from common_code import kwq, util, file_names, vector_utils
import time
from multiprocessing import Process, Queue, Manager
import threading
topic_subscribe = 'upload_to_vector_db'
group_id = "vision_v3_uploaders"
def get_score_file(f):
return f.key().decode().strip('"')
def watch_and_upload(rank, health_queue):
start_time = time.time()
client_suffix = '_'+str(rank)
client_id = 'uploader_to_vector_db'+client_suffix
queue = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'earliest')
queue.subscribe([topic_subscribe])
last_heartbeat = time.time()
while True:
current_time = time.time()
# # Exit after 5 minutes
# if (current_time - start_time) > 60*5:
# print('###########################')
# print('######### EXITING #########')
# print('###########################')
# return
# Send heartbeat every 10 seconds
if current_time - last_heartbeat > 5:
try:
health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'alive'})
last_heartbeat = current_time
except:
# Queue might be full, continue working
pass
result = queue.poll(timeout=0.05)
if result is not None:
score_this = get_score_file(result)
try:
upload_vectors_to_db(score_this)
print('#' + str(rank) + ' Uploaded to ' + score_this)
# Send success heartbeat
try:
health_queue.put({'rank': rank, 'timestamp': current_time, 'status': 'working'})
except:
pass
except Exception as e:
print('Failed ' + score_this + ':' + str(e))
queue.commit()
def health_monitor(processes, health_queues, num_workers):
"""Monitor process health and restart hung processes"""
last_heartbeat = {}
HEALTH_TIMEOUT = 15 # Consider process hung if no heartbeat for 60 seconds
while True:
time.sleep(5) # Check every 5 seconds
current_time = time.time()
for rank in range(num_workers):
# Drain the health queue for this process
latest_heartbeat = None
try:
while True:
try:
heartbeat = health_queues[rank].get_nowait()
if heartbeat['rank'] == rank:
latest_heartbeat = heartbeat['timestamp']
except:
break
except:
pass
# Update last known heartbeat
if latest_heartbeat:
last_heartbeat[rank] = latest_heartbeat
# Check if process is alive and responsive
process = processes[rank]
last_seen = last_heartbeat.get(rank, 0)
if not process.is_alive():
print(f"Process {rank} has died. Restarting...")
processes[rank] = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
processes[rank].start()
last_heartbeat[rank] = current_time
elif current_time - last_seen > HEALTH_TIMEOUT:
print(f"Process {rank} appears hung (no heartbeat for {current_time - last_seen:.1f}s). Restarting...")
process.terminate()
process.join(timeout=5)
if process.is_alive():
process.kill()
process.join()
# Create new process
processes[rank] = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
processes[rank].start()
last_heartbeat[rank] = current_time
# Main execution
if __name__ == "__main__":
num_workers = 4
# Create health queues for each worker
health_queues = [Queue(maxsize=10) for _ in range(num_workers)]
# Create and start worker processes
processes = []
for rank in range(num_workers):
p = Process(target=watch_and_upload, args=(rank, health_queues[rank]))
p.start()
processes.append(p)
# Start health monitor in a separate thread
monitor_thread = threading.Thread(
target=health_monitor,
args=(processes, health_queues, num_workers),
daemon=True
)
monitor_thread.start()
try:
# Wait for all processes to complete (they will exit after 5 minutes)
for p in processes:
p.join()
print("All worker processes completed")
except KeyboardInterrupt:
print("Shutting down...")
for p in processes:
p.terminate()
p.join(timeout=5)
if p.is_alive():
p.kill()
p.join()