Skip to content

V
vidai
Commit d370350d authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

Implement advanced job scheduling system with VRAM-aware worker selection, job... 

Implement advanced job scheduling system with VRAM-aware worker selection, job queuing, RunPod integration, model management, and real-time notifications

- Add VRAM requirement determination for models
- Implement intelligent worker selection based on VRAM availability and weights
- Add job queuing when no workers available with automatic retry
- Integrate RunPod pod creation and management for scaling
- Implement model loading/unloading with reference counting
- Add file transfer support for remote workers (shared storage + websocket)
- Enable concurrent job processing with VRAM tracking
- Create job results page with detailed output display
- Add real-time job completion notifications with result links
- Update history page with live progress updates and result links
- Fix async handling in cluster master websocket communication
- Add database schema updates for job tracking
parent 4bc78616
Expand all Show whitespace changes
Inline Side-by-side
Showing with 921 additions and 119 deletions
templates/history.html
View file @ d370350d
......@@ -20,7 +20,125 @@
.status-failed { background: #fee2e2; color: #dc2626; }
.job-tokens { font-weight: 600; color: #667eea; text-align: center; }
.no-jobs { text-align: center; padding: 3rem; color: #6b7280; }
.job-progress { color: #64748b; font-size: 0.8rem; }
.spinner { display: inline-block; width: 12px; height: 12px; border: 2px solid #f3f3f3; border-top: 2px solid #667eea; border-radius: 50%; animation: spin 1s linear infinite; margin-right: 0.5rem; }
.view-result-link { color: #667eea; text-decoration: none; font-size: 0.9rem; font-weight: 500; }
.view-result-link:hover { text-decoration: underline; }
@keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }
</style>
<script>
let lastUpdate = Date.now();
let lastCompletedJobs = new Set(); // Track jobs that were already completed
function updateJobStatuses() {
fetch('/api/job_status_updates?since=' + lastUpdate)
.then(response => response.json())
.then(data => {
if (data.updates && data.updates.length > 0) {
data.updates.forEach(update => {
updateJobRow(update);
// Check if this is a newly completed job
if (update.status === 'completed' && !lastCompletedJobs.has(update.id)) {
showJobCompletionNotification(update.id);
lastCompletedJobs.add(update.id);
}
});
lastUpdate = Date.now();
}
})
.catch(error => {
console.log('Error updating job statuses:', error);
});
}
function showJobCompletionNotification(jobId) {
const notificationContainer = document.getElementById('notificationContainer');
if (!notificationContainer) return;
const notification = document.createElement('div');
notification.className = 'notification success';
notification.innerHTML = `
<span class="notification-close" onclick="closeNotification(this)">&times;</span>
<strong>Job Completed!</strong><br>
Your analysis job has finished. <a href="/job_result/${jobId}" style="color: #065f46; text-decoration: underline;">View Results</a>
`;
notificationContainer.appendChild(notification);
// Auto-hide after 10 seconds
setTimeout(() => {
notification.classList.add('fade-out');
setTimeout(() => {
notification.remove();
}, 300);
}, 10000);
}
function updateJobRow(jobUpdate) {
const jobRow = document.querySelector(`[data-job-id="${jobUpdate.id}"]`);
if (jobRow) {
// Update status
const statusElement = jobRow.querySelector('.job-status');
if (statusElement) {
statusElement.className = `job-status status-${jobUpdate.status}`;
statusElement.textContent = jobUpdate.status.charAt(0).toUpperCase() + jobUpdate.status.slice(1);
// Add spinner for processing jobs
if (jobUpdate.status === 'processing') {
if (!statusElement.querySelector('.spinner')) {
statusElement.innerHTML = '<div class="spinner"></div>' + statusElement.textContent;
}
} else {
const spinner = statusElement.querySelector('.spinner');
if (spinner) {
spinner.remove();
}
}
}
// Update tokens if completed
if (jobUpdate.status === 'completed' && jobUpdate.used_tokens) {
const tokensElement = jobRow.querySelector('.job-tokens');
if (tokensElement) {
tokensElement.textContent = jobUpdate.used_tokens;
}
}
// Add progress info for processing jobs
if (jobUpdate.status === 'processing' && jobUpdate.result) {
let progressElement = jobRow.querySelector('.job-progress');
if (!progressElement) {
progressElement = document.createElement('div');
progressElement.className = 'job-progress';
jobRow.appendChild(progressElement);
}
progressElement.textContent = jobUpdate.result.status || 'Processing...';
}
}
}
// Initialize completed jobs tracking
function initializeCompletedJobs() {
const jobRows = document.querySelectorAll('[data-job-id]');
jobRows.forEach(row => {
const jobId = parseInt(row.getAttribute('data-job-id'));
const statusElement = row.querySelector('.job-status');
if (statusElement && statusElement.classList.contains('status-completed')) {
lastCompletedJobs.add(jobId);
}
});
}
// Update every 5 seconds
setInterval(updateJobStatuses, 5000);
// Initial update
document.addEventListener('DOMContentLoaded', function() {
initializeCompletedJobs();
updateJobStatuses();
});
</script>
{% endblock %}
{% block content %}
......@@ -30,14 +148,22 @@
<h2><i class="fas fa-history"></i> Job History</h2>
</div>
{% for job in queue_items %}
<div class="job-row">
<div class="job-row" data-job-id="{{ job.id }}">
<div class="job-type">{{ job.request_type.title() }}</div>
<div class="job-data" title="{{ job.data.get('prompt', job.data.get('description', 'N/A')) }}">
{{ job.data.get('prompt', job.data.get('description', 'N/A'))[:50] }}{% if job.data.get('prompt', job.data.get('description', 'N/A'))|length > 50 %}...{% endif %}
</div>
<div class="job-time">{{ job.created_at[:19] }}</div>
<span class="job-status status-{{ job.status }}">{{ job.status.title() }}</span>
<span class="job-status status-{{ job.status }}">
{% if job.status == 'processing' %}<div class="spinner"></div>{% endif %}
{{ job.status.title() }}
</span>
<div class="job-tokens">{{ job.used_tokens or 0 }}</div>
{% if job.status == 'completed' %}
<a href="/job_result/{{ job.id }}" class="view-result-link">View Result</a>
{% elif job.status == 'processing' and job.result %}
<div class="job-progress">{{ job.result.get('status', 'Processing...') }}</div>
{% endif %}
</div>
{% endfor %}
{% if not queue_items %}
......
templates/job_result.html 0 → 100644
View file @ d370350d
{% extends "base.html" %}
{% block title %}Job Result - VidAI{% endblock %}
{% block head %}
<style>
.container { max-width: 1200px; margin: 2rem auto; padding: 0 2rem; }
.result-card { background: white; border-radius: 12px; box-shadow: 0 2px 10px rgba(0,0,0,0.05); padding: 2rem; margin-bottom: 2rem; }
.result-header { display: flex; justify-content: space-between; align-items: center; margin-bottom: 1.5rem; }
.result-title { margin: 0; color: #1e293b; }
.result-meta { color: #64748b; font-size: 0.9rem; }
.result-content { line-height: 1.6; white-space: pre-wrap; }
.back-link { color: #667eea; text-decoration: none; }
.back-link:hover { text-decoration: underline; }
.result-actions { display: flex; gap: 1rem; }
.btn { padding: 0.5rem 1rem; border: none; border-radius: 6px; text-decoration: none; font-weight: 500; cursor: pointer; }
.btn-primary { background: #667eea; color: white; }
.btn-primary:hover { background: #5a67d8; }
.btn-secondary { background: #f1f5f9; color: #475569; }
.btn-secondary:hover { background: #e2e8f0; }
</style>
{% endblock %}
{% block content %}
<div class="container">
<div class="result-card">
<div class="result-header">
<div>
<h1 class="result-title">{{ job.request_type.title() }} Result</h1>
<div class="result-meta">
Completed on {{ job.completed_at[:19] if job.completed_at else 'Unknown' }} •
Used {{ job.used_tokens or 0 }} tokens
</div>
</div>
<div class="result-actions">
<a href="/history" class="btn btn-secondary">← Back to History</a>
<button onclick="window.print()" class="btn btn-primary">Print Result</button>
</div>
</div>
<div class="result-content">
{% if job.result %}
{% if job.result.get('result') %}
{{ job.result.result }}
{% else %}
{{ job.result | tojson(indent=2) }}
{% endif %}
{% else %}
<em>No result data available</em>
{% endif %}
</div>
</div>
{% if job.data %}
<div class="result-card">
<h3>Job Details</h3>
<div class="result-meta">
<strong>Prompt:</strong> {{ job.data.get('prompt', 'N/A') }}<br>
<strong>Model:</strong> {{ job.data.get('model_path', 'N/A') }}<br>
<strong>Interval:</strong> {{ job.data.get('interval', 'N/A') }} seconds<br>
<strong>Created:</strong> {{ job.created_at[:19] }}
</div>
</div>
{% endif %}
</div>
{% endblock %}
\ No newline at end of file
vidai/backend.py
View file @ d370350d
......@@ -32,16 +32,24 @@ pending_results = {} # msg_id -> result message
def handle_web_message(message: Message) -> Message:
"""Handle messages from web interface."""
if message.msg_type == 'analyze_request':
backend = get_analysis_backend()
worker_key = f'analysis_{backend}'
if worker_key in worker_sockets:
# Forward to worker
worker_sockets[worker_key].sendall(
f'{{"msg_type": "{message.msg_type}", "msg_id": "{message.msg_id}", "data": {message.data}}}\n'.encode('utf-8')
)
from .cluster_master import cluster_master
import asyncio
# Use advanced job scheduling
try:
# Run async function in sync context
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
job_id = loop.run_until_complete(cluster_master.assign_job_with_model('analysis', message.data))
loop.close()
if job_id:
# Job assigned, will respond asynchronously
return None # No immediate response
else:
return Message('error', message.msg_id, {'error': f'Worker {worker_key} not available'})
return Message('error', message.msg_id, {'error': 'No suitable worker available'})
except Exception as e:
return Message('error', message.msg_id, {'error': f'Job scheduling failed: {str(e)}'})
elif message.msg_type == 'train_request':
backend = get_training_backend()
worker_key = f'training_{backend}'
......
vidai/cluster_client.py
View file @ d370350d
......@@ -218,6 +218,18 @@ class ClusterClient:
elif msg_type == 'job_assignment':
await self._handle_job_assignment(message)
elif msg_type == 'job_file_shared':
await self._handle_job_file_shared(message)
elif msg_type == 'job_file_transfer_start':
await self._handle_job_file_transfer_start(message)
elif msg_type == 'job_file_chunk':
await self._handle_job_file_chunk(message)
elif msg_type == 'job_file_transfer_complete':
await self._handle_job_file_transfer_complete(message)
elif msg_type == 'receive_file':
await self._handle_receive_file(message)
......@@ -281,12 +293,54 @@ class ClusterClient:
"""Handle job assignment from master."""
job_id = message.get('job_id')
job_data = message.get('job_data', {})
# Process job locally and send result back
# This is a placeholder - actual implementation would depend on job type
# Extract job parameters
request_type = job_data.get('request_type', 'analyze')
model_path = job_data.get('model_path', 'Qwen/Qwen2.5-VL-7B-Instruct')
media_path = job_data.get('local_path')
prompt = job_data.get('prompt', 'Describe this image.')
interval = job_data.get('interval', 10)
# Forward to appropriate local worker
worker_type = f'{request_type}_cuda' # Assume CUDA for now
if worker_type in self.local_processes:
# Send job to local worker process
import json
job_message = {
'msg_type': f'{request_type}_request',
'msg_id': job_id,
'data': {
'model_path': model_path,
'local_path': media_path,
'prompt': prompt,
'interval': interval
}
}
# Send to worker via socket or other mechanism
# For now, assume workers listen on sockets
try:
import socket
worker_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
worker_socket.connect(f'/tmp/vidai_{worker_type}.sock')
worker_socket.sendall(json.dumps(job_message).encode('utf-8'))
worker_socket.close()
# Wait for result (simplified)
await asyncio.sleep(1) # Placeholder
except Exception as e:
print(f"Failed to send job to local worker: {e}")
await self._send_message({
'type': 'job_result',
'job_id': job_id,
'result': {'status': 'completed', 'data': 'placeholder result'}
'result': {'status': 'failed', 'error': str(e)}
})
else:
await self._send_message({
'type': 'job_result',
'job_id': job_id,
'result': {'status': 'failed', 'error': f'No local {worker_type} worker available'}
})
async def _handle_receive_file(self, message: Dict[str, Any]) -> None:
......@@ -481,6 +535,82 @@ class ClusterClient:
except Exception as e:
print(f"Error handling shared model file {shared_file_path}: {e}")
async def _handle_job_file_shared(self, message: Dict[str, Any]) -> None:
"""Handle job file available in shared directory."""
job_id = message.get('job_id')
shared_file_path = message.get('shared_file_path')
original_path = message.get('original_path')
if not self.shared_dir:
print(f"Received shared job file message but no shared directory configured: {shared_file_path}")
return
try:
# Copy shared file to local temp location
import shutil
import tempfile
local_temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(shared_file_path)[1])
shutil.copy2(shared_file_path, local_temp_file.name)
local_temp_file.close()
print(f"Job file for {job_id} copied from shared directory: {shared_file_path} -> {local_temp_file.name}")
# Update job data with local path
# This would need to be stored and used when processing the job
# For now, just log
print(f"Job {job_id} file ready at: {local_temp_file.name}")
except Exception as e:
print(f"Error handling shared job file {shared_file_path}: {e}")
async def _handle_job_file_transfer_start(self, message: Dict[str, Any]) -> None:
"""Handle start of job file transfer."""
job_id = message.get('job_id')
file_path = message.get('file_path')
total_size = message.get('total_size', 0)
# Initialize file transfer
self._current_job_file_transfer = {
'job_id': job_id,
'file_path': file_path,
'total_size': total_size,
'received_data': b''
}
print(f"Starting job file transfer for {job_id}: {file_path} ({total_size} bytes)")
async def _handle_job_file_chunk(self, message: Dict[str, Any]) -> None:
"""Handle job file data chunk."""
if not hasattr(self, '_current_job_file_transfer'):
return
chunk_hex = message.get('data', '')
chunk_data = bytes.fromhex(chunk_hex)
self._current_job_file_transfer['received_data'] += chunk_data
async def _handle_job_file_transfer_complete(self, message: Dict[str, Any]) -> None:
"""Handle completion of job file transfer."""
if not hasattr(self, '_current_job_file_transfer'):
return
job_id = self._current_job_file_transfer['job_id']
file_path = self._current_job_file_transfer['file_path']
received_data = self._current_job_file_transfer['received_data']
expected_size = self._current_job_file_transfer['total_size']
if len(received_data) == expected_size:
# Save the file
import tempfile
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file_path)[1])
temp_file.write(received_data)
temp_file.close()
print(f"Job file for {job_id} saved to: {temp_file.name}")
# Clean up
delattr(self, '_current_job_file_transfer')
else:
print(f"Job file transfer failed: received {len(received_data)} bytes, expected {expected_size}")
async def run(self) -> None:
"""Main client loop with reconnection."""
reconnect = True
......
vidai/cluster_master.py
View file @ d370350d
This diff is collapsed.
vidai/database.py
View file @ d370350d
......@@ -254,6 +254,7 @@ def init_db(conn) -> None:
estimated_time INT,
estimated_tokens INT DEFAULT 0,
used_tokens INT DEFAULT 0,
job_id VARCHAR(100),
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
''')
......@@ -274,7 +275,7 @@ def init_db(conn) -> None:
estimated_time INTEGER,
estimated_tokens INTEGER DEFAULT 0,
used_tokens INTEGER DEFAULT 0,
FOREIGN KEY (user_id) REFERENCES users (id)
job_id TEXT
)
''')
......@@ -1015,7 +1016,7 @@ def delete_token_package(package_id: int) -> bool:
# Queue management functions
def add_to_queue(user_id: int, request_type: str, data: dict, priority: int = 0) -> int:
def add_to_queue(user_id: int, request_type: str, data: dict, priority: int = 0, job_id: str = None) -> int:
"""Add request to processing queue."""
import json
......@@ -1025,9 +1026,9 @@ def add_to_queue(user_id: int, request_type: str, data: dict, priority: int = 0)
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute('''
INSERT INTO processing_queue (user_id, request_type, data, priority, estimated_tokens)
VALUES (?, ?, ?, ?, ?)
''', (user_id, request_type, json.dumps(data), priority, estimated_tokens))
INSERT INTO processing_queue (user_id, request_type, data, priority, estimated_tokens, job_id)
VALUES (?, ?, ?, ?, ?, ?)
''', (user_id, request_type, json.dumps(data), priority, estimated_tokens, job_id))
queue_id = cursor.lastrowid
conn.commit()
conn.close()
......@@ -1084,6 +1085,21 @@ def get_queue_status(queue_id: int) -> Optional[Dict[str, Any]]:
return None
def get_queue_by_job_id(job_id: str) -> Optional[Dict[str, Any]]:
"""Get queue item by cluster job_id."""
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute('SELECT * FROM processing_queue WHERE job_id = ?', (job_id,))
row = cursor.fetchone()
conn.close()
if row:
item = dict(row)
item['data'] = json.loads(item['data']) if item['data'] else None
item['result'] = json.loads(item['result']) if item['result'] else None
return item
return None
def get_pending_queue_items() -> List[Dict[str, Any]]:
"""Get pending queue items ordered by priority and creation time."""
conn = get_db_connection()
......@@ -1103,7 +1119,7 @@ def get_pending_queue_items() -> List[Dict[str, Any]]:
return items
def update_queue_status(queue_id: int, status: str, result: dict = None, error: str = None, used_tokens: int = None) -> bool:
def update_queue_status(queue_id: int, status: str, result: dict = None, error: str = None, used_tokens: int = None, job_id: str = None) -> bool:
"""Update queue item status."""
import json
conn = get_db_connection()
......@@ -1131,6 +1147,10 @@ def update_queue_status(queue_id: int, status: str, result: dict = None, error:
update_fields.append('used_tokens = ?')
params.append(used_tokens)
if job_id is not None:
update_fields.append('job_id = ?')
params.append(job_id)
params.append(queue_id)
query = f'UPDATE processing_queue SET {", ".join(update_fields)} WHERE id = ?'
......
vidai/models.py
View file @ d370350d
......@@ -240,6 +240,67 @@ def get_model(model_path: str, model_type: str = None, **kwargs) -> BaseModel:
return model
def get_model_vram_requirement(model_path: str) -> int:
"""Get VRAM requirement in GB for a model."""
model_path_lower = model_path.lower()
# Known model VRAM requirements (in GB)
vram_requirements = {
# Qwen models
'qwen/qwen2.5-vl-7b-instruct': 16,
'qwen/qwen2-vl-7b-instruct': 16,
'qwen/qwen-vl-chat': 16,
'qwen/qwen2.5-vl-3b-instruct': 8,
'qwen/qwen2.5-vl-72b-instruct': 144,
# LLaMA models
'meta-llama/llama-3.1-8b-instruct': 16,
'meta-llama/llama-3.1-70b-instruct': 140,
'meta-llama/llama-3.1-405b-instruct': 800,
'meta-llama/llama-3-8b-instruct': 16,
'meta-llama/llama-3-70b-instruct': 140,
# Mistral models
'mistralai/mistral-7b-instruct': 14,
'mistralai/mixtral-8x7b-instruct': 56,
'mistralai/mistral-large': 120,
# Other models
'microsoft/wizardlm-2-8x22b': 180,
'databricks/dbrx-instruct': 120,
}
# Check for exact matches first
if model_path in vram_requirements:
return vram_requirements[model_path]
# Check for partial matches
for model_key, vram in vram_requirements.items():
if model_key in model_path_lower or model_path_lower in model_key:
return vram
# Estimate based on model name patterns
if '72b' in model_path_lower or '70b' in model_path_lower:
return 140
elif '13b' in model_path_lower or '14b' in model_path_lower:
return 28
elif '7b' in model_path_lower or '8b' in model_path_lower:
return 16
elif '3b' in model_path_lower:
return 8
elif '1b' in model_path_lower:
return 4
elif '405b' in model_path_lower:
return 800
elif 'mixtral' in model_path_lower or '8x7b' in model_path_lower:
return 56
elif '8x22b' in model_path_lower:
return 180
# Default fallback
return 16 # Assume 16GB for unknown models
def unload_all_models() -> None:
"""Unload all cached models."""
for model in _model_cache.values():
......
vidai/queue.py
View file @ d370350d
......@@ -120,20 +120,30 @@ class QueueManager:
def _execute_local_or_distributed_job(self, job: Dict[str, Any]) -> None:
"""Execute job using local workers or distributed cluster."""
import asyncio
from .cluster_master import cluster_master
# Determine process type
process_type = job['request_type'] # 'analyze' or 'train'
# Try to get best available worker
worker_key = cluster_master.get_best_worker(process_type)
if worker_key:
# Send to distributed worker
self._send_to_distributed_worker(job, worker_key)
# Use advanced job scheduling
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
job_id = loop.run_until_complete(cluster_master.assign_job_with_model(process_type, job['data']))
loop.close()
if job_id:
# Job assigned successfully, mark as processing and store job_id
from .database import update_queue_status
update_queue_status(job['id'], 'processing', {'job_id': job_id, 'status': 'Assigned to worker'}, job_id=job_id)
else:
# Fall back to local processing
self._execute_local_job(job)
# No worker available, leave in queue for retry
print(f"No worker available for job {job['id']}, will retry later")
# Don't update status, leave as 'queued'
except Exception as e:
from .database import update_queue_status
update_queue_status(job['id'], 'failed', error_message=str(e))
def _send_to_distributed_worker(self, job: Dict[str, Any], worker_key: str) -> None:
"""Send job to distributed worker."""
......
vidai/web.py
View file @ d370350d
......@@ -288,15 +288,14 @@ def analyze():
'interval': interval,
'user_id': user['id']
}
msg_id = send_to_backend('analyze_request', data)
result_data = get_result(msg_id)
if 'data' in result_data:
result = result_data['data'].get('result', 'Analysis completed')
# Deduct tokens (skip for admin users)
if user.get('role') != 'admin':
update_user_tokens(user['id'], -10)
else:
result = result_data.get('error', 'Error')
# Submit job to queue system
from .queue import queue_manager
job_id = queue_manager.submit_job(user['id'], 'analyze', data)
# For immediate response, we could poll, but for now redirect to history
flash('Analysis job submitted successfully! Check your history for results.', 'success')
return redirect(url_for('history'))
return render_template('analyze.html',
user=user,
......@@ -322,6 +321,70 @@ def history():
active_page='history')
@app.route('/job_result/<int:job_id>')
@login_required
def job_result(job_id):
"""Job result page."""
user = get_current_user_session()
# Get the job details
job = get_queue_status(job_id)
# Check if job belongs to user and is completed
if not job or job['user_id'] != user['id']:
flash('Job not found or access denied', 'error')
return redirect(url_for('history'))
if job['status'] != 'completed':
flash('Job is not completed yet', 'error')
return redirect(url_for('history'))
return render_template('job_result.html',
user=user,
job=job,
active_page='history')
@app.route('/api/job_status_updates')
@login_required
def api_job_status_updates():
"""API endpoint for job status updates."""
user = get_current_user_session()
since = request.args.get('since', '0')
try:
since_timestamp = int(since) / 1000.0 # Convert from milliseconds to seconds
except ValueError:
since_timestamp = 0
# Get updated jobs since the given timestamp
queue_items = get_user_queue_items(user['id'])
updates = []
for job in queue_items:
# Convert job timestamps to comparable format
job_time = job.get('created_at')
if isinstance(job_time, str):
# Parse timestamp string
from datetime import datetime
try:
job_timestamp = datetime.fromisoformat(job_time.replace('Z', '+00:00')).timestamp()
except:
job_timestamp = 0
else:
job_timestamp = job_time or 0
if job_timestamp > since_timestamp:
updates.append({
'id': job['id'],
'status': job['status'],
'used_tokens': job.get('used_tokens', 0),
'result': job.get('result', {})
})
return {'updates': updates}
@app.route('/update_settings', methods=['POST'])
@login_required
def update_settings():
......
vidai/worker_analysis.py
View file @ d370350d
......@@ -44,6 +44,10 @@ if torch.cuda.is_available():
else:
max_gpu = min_gpu = 0
# Model management for workers
loaded_models = {} # model_path -> (model_instance, ref_count)
current_model_path = None # Currently active model
# Set OpenCV to smaller GPU if available
try:
if cv2 and hasattr(cv2, 'cuda'):
......@@ -93,6 +97,49 @@ def extract_frames(video_path, interval=10, optimize=False):
def is_video(file_path):
return file_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))
def get_or_load_model(model_path: str):
"""Get a model instance, loading it if necessary."""
global current_model_path
if model_path in loaded_models:
# Model already loaded, increment ref count
model, ref_count = loaded_models[model_path]
loaded_models[model_path] = (model, ref_count + 1)
current_model_path = model_path
return model
# Check if we need to unload current model
if current_model_path and current_model_path != model_path:
# Check if current model is still referenced by other jobs
if loaded_models[current_model_path][1] <= 1:
# Only this job is using it, unload it
unload_model(current_model_path)
# Note: We don't change current_model_path yet
# Load new model
model = get_model(model_path)
loaded_models[model_path] = (model, 1)
current_model_path = model_path
return model
def unload_model(model_path: str):
"""Unload a model if it's no longer referenced."""
if model_path in loaded_models:
model, ref_count = loaded_models[model_path]
if ref_count <= 1:
# No more references, unload
model.unload_model()
del loaded_models[model_path]
if current_model_path == model_path:
current_model_path = None
else:
# Decrement ref count
loaded_models[model_path] = (model, ref_count - 1)
def release_model(model_path: str):
"""Release a reference to a model (called when job completes)."""
unload_model(model_path) # This will decrement ref count
def analyze_single_image(image_path, prompt, model):
"""Analyze a single image using the dynamic model."""
messages = [
......@@ -111,8 +158,8 @@ def analyze_media(media_path, prompt, model_path, interval=10):
"""Analyze media using dynamic model loading."""
torch.cuda.empty_cache()
# Get model dynamically
model = get_model(model_path, model_type=None) # Auto-detect type
# Get model with reference counting
model = get_or_load_model(model_path)
# Get system prompt
try:
......@@ -180,8 +227,15 @@ def worker_process(backend_type: str):
interval = data.get('interval', 10)
result = analyze_media(media_path, prompt, model_path, interval)
# Release model reference (don't unload yet, per requirements)
release_model(model_path)
# Send result back
response = Message('analyze_response', message.msg_id, {'result': result})
comm.send_message(response)
# If in cluster mode, also notify cluster master
# This would be handled by the cluster client receiving the result
time.sleep(0.1)
except Exception as e:
print(f"Worker error: {e}")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment