445 lines
17 KiB
Python
445 lines
17 KiB
Python
|
"""
|
||
|
|
Job Queue Repository
|
||
|
|
Repository for training job queue operations
|
||
|
|
"""
|
||
|
|
|
||
|
|
from typing import Optional, List, Dict, Any
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
from sqlalchemy import select, and_, text, desc
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
import structlog
|
||
|
|
|
||
|
|
from .base import TrainingBaseRepository
|
||
|
|
from app.models.training import TrainingJobQueue
|
||
|
|
from shared.database.exceptions import DatabaseError, ValidationError
|
||
|
|
|
||
|
|
logger = structlog.get_logger()
|
||
|
|
|
||
|
|
|
||
|
|
class JobQueueRepository(TrainingBaseRepository):
|
||
|
|
"""Repository for training job queue operations"""
|
||
|
|
|
||
|
|
def __init__(self, session: AsyncSession, cache_ttl: Optional[int] = 60):
|
||
|
|
# Job queue changes frequently, very short cache time (1 minute)
|
||
|
|
super().__init__(TrainingJobQueue, session, cache_ttl)
|
||
|
|
|
||
|
|
async def enqueue_job(self, job_data: Dict[str, Any]) -> TrainingJobQueue:
|
||
|
|
"""Add a job to the training queue"""
|
||
|
|
try:
|
||
|
|
# Validate job data
|
||
|
|
validation_result = self._validate_training_data(
|
||
|
|
job_data,
|
||
|
|
["job_id", "tenant_id", "job_type"]
|
||
|
|
)
|
||
|
|
|
||
|
|
if not validation_result["is_valid"]:
|
||
|
|
raise ValidationError(f"Invalid job data: {validation_result['errors']}")
|
||
|
|
|
||
|
|
# Set default values
|
||
|
|
if "priority" not in job_data:
|
||
|
|
job_data["priority"] = 1
|
||
|
|
if "status" not in job_data:
|
||
|
|
job_data["status"] = "queued"
|
||
|
|
if "max_retries" not in job_data:
|
||
|
|
job_data["max_retries"] = 3
|
||
|
|
|
||
|
|
# Create queue entry
|
||
|
|
queued_job = await self.create(job_data)
|
||
|
|
|
||
|
|
logger.info("Job enqueued",
|
||
|
|
job_id=queued_job.job_id,
|
||
|
|
tenant_id=queued_job.tenant_id,
|
||
|
|
job_type=queued_job.job_type,
|
||
|
|
priority=queued_job.priority)
|
||
|
|
|
||
|
|
return queued_job
|
||
|
|
|
||
|
|
except ValidationError:
|
||
|
|
raise
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to enqueue job",
|
||
|
|
job_id=job_data.get("job_id"),
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to enqueue job: {str(e)}")
|
||
|
|
|
||
|
|
async def get_next_job(self, job_types: List[str] = None) -> Optional[TrainingJobQueue]:
|
||
|
|
"""Get the next job to process from the queue"""
|
||
|
|
try:
|
||
|
|
# Build filters for job types if specified
|
||
|
|
filters = {"status": "queued"}
|
||
|
|
|
||
|
|
if job_types:
|
||
|
|
# For multiple job types, we need to use raw SQL
|
||
|
|
job_types_str = "', '".join(job_types)
|
||
|
|
query_text = f"""
|
||
|
|
SELECT * FROM training_job_queue
|
||
|
|
WHERE status = 'queued'
|
||
|
|
AND job_type IN ('{job_types_str}')
|
||
|
|
AND (scheduled_at IS NULL OR scheduled_at <= :now)
|
||
|
|
ORDER BY priority DESC, created_at ASC
|
||
|
|
LIMIT 1
|
||
|
|
"""
|
||
|
|
|
||
|
|
result = await self.session.execute(text(query_text), {"now": datetime.now()})
|
||
|
|
row = result.fetchone()
|
||
|
|
|
||
|
|
if row:
|
||
|
|
record_dict = dict(row._mapping)
|
||
|
|
return self.model(**record_dict)
|
||
|
|
return None
|
||
|
|
else:
|
||
|
|
# Simple case - get any queued job
|
||
|
|
jobs = await self.get_multi(
|
||
|
|
filters=filters,
|
||
|
|
limit=1,
|
||
|
|
order_by="priority",
|
||
|
|
order_desc=True
|
||
|
|
)
|
||
|
|
return jobs[0] if jobs else None
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to get next job from queue",
|
||
|
|
job_types=job_types,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to get next job: {str(e)}")
|
||
|
|
|
||
|
|
async def start_job(self, job_id: str) -> Optional[TrainingJobQueue]:
|
||
|
|
"""Mark a job as started"""
|
||
|
|
try:
|
||
|
|
job = await self.get_by_job_id(job_id)
|
||
|
|
if not job:
|
||
|
|
logger.error(f"Job not found in queue: {job_id}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
if job.status != "queued":
|
||
|
|
logger.warning(f"Job {job_id} is not queued (status: {job.status})")
|
||
|
|
return job
|
||
|
|
|
||
|
|
updated_job = await self.update(job.id, {
|
||
|
|
"status": "running",
|
||
|
|
"started_at": datetime.now(),
|
||
|
|
"updated_at": datetime.now()
|
||
|
|
})
|
||
|
|
|
||
|
|
logger.info("Job started",
|
||
|
|
job_id=job_id,
|
||
|
|
job_type=job.job_type)
|
||
|
|
|
||
|
|
return updated_job
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to start job",
|
||
|
|
job_id=job_id,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to start job: {str(e)}")
|
||
|
|
|
||
|
|
async def complete_job(self, job_id: str) -> Optional[TrainingJobQueue]:
|
||
|
|
"""Mark a job as completed"""
|
||
|
|
try:
|
||
|
|
job = await self.get_by_job_id(job_id)
|
||
|
|
if not job:
|
||
|
|
logger.error(f"Job not found in queue: {job_id}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
updated_job = await self.update(job.id, {
|
||
|
|
"status": "completed",
|
||
|
|
"updated_at": datetime.now()
|
||
|
|
})
|
||
|
|
|
||
|
|
logger.info("Job completed",
|
||
|
|
job_id=job_id,
|
||
|
|
job_type=job.job_type if job else "unknown")
|
||
|
|
|
||
|
|
return updated_job
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to complete job",
|
||
|
|
job_id=job_id,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to complete job: {str(e)}")
|
||
|
|
|
||
|
|
async def fail_job(self, job_id: str, error_message: str = None) -> Optional[TrainingJobQueue]:
|
||
|
|
"""Mark a job as failed and handle retries"""
|
||
|
|
try:
|
||
|
|
job = await self.get_by_job_id(job_id)
|
||
|
|
if not job:
|
||
|
|
logger.error(f"Job not found in queue: {job_id}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Increment retry count
|
||
|
|
new_retry_count = job.retry_count + 1
|
||
|
|
|
||
|
|
# Check if we should retry
|
||
|
|
if new_retry_count < job.max_retries:
|
||
|
|
# Reset to queued for retry
|
||
|
|
updated_job = await self.update(job.id, {
|
||
|
|
"status": "queued",
|
||
|
|
"retry_count": new_retry_count,
|
||
|
|
"updated_at": datetime.now(),
|
||
|
|
"started_at": None # Reset started_at for retry
|
||
|
|
})
|
||
|
|
|
||
|
|
logger.info("Job failed, queued for retry",
|
||
|
|
job_id=job_id,
|
||
|
|
retry_count=new_retry_count,
|
||
|
|
max_retries=job.max_retries)
|
||
|
|
else:
|
||
|
|
# Mark as permanently failed
|
||
|
|
updated_job = await self.update(job.id, {
|
||
|
|
"status": "failed",
|
||
|
|
"retry_count": new_retry_count,
|
||
|
|
"updated_at": datetime.now()
|
||
|
|
})
|
||
|
|
|
||
|
|
logger.error("Job permanently failed",
|
||
|
|
job_id=job_id,
|
||
|
|
retry_count=new_retry_count,
|
||
|
|
error_message=error_message)
|
||
|
|
|
||
|
|
return updated_job
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to handle job failure",
|
||
|
|
job_id=job_id,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to handle job failure: {str(e)}")
|
||
|
|
|
||
|
|
async def cancel_job(self, job_id: str, cancelled_by: str = None) -> Optional[TrainingJobQueue]:
|
||
|
|
"""Cancel a job"""
|
||
|
|
try:
|
||
|
|
job = await self.get_by_job_id(job_id)
|
||
|
|
if not job:
|
||
|
|
logger.error(f"Job not found in queue: {job_id}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
if job.status in ["completed", "failed"]:
|
||
|
|
logger.warning(f"Cannot cancel job {job_id} with status {job.status}")
|
||
|
|
return job
|
||
|
|
|
||
|
|
updated_job = await self.update(job.id, {
|
||
|
|
"status": "cancelled",
|
||
|
|
"cancelled_by": cancelled_by,
|
||
|
|
"updated_at": datetime.now()
|
||
|
|
})
|
||
|
|
|
||
|
|
logger.info("Job cancelled",
|
||
|
|
job_id=job_id,
|
||
|
|
cancelled_by=cancelled_by)
|
||
|
|
|
||
|
|
return updated_job
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to cancel job",
|
||
|
|
job_id=job_id,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to cancel job: {str(e)}")
|
||
|
|
|
||
|
|
async def get_queue_status(self, tenant_id: str = None) -> Dict[str, Any]:
|
||
|
|
"""Get queue status and statistics"""
|
||
|
|
try:
|
||
|
|
base_filters = {}
|
||
|
|
if tenant_id:
|
||
|
|
base_filters["tenant_id"] = tenant_id
|
||
|
|
|
||
|
|
# Get counts by status
|
||
|
|
queued_jobs = await self.count(filters={**base_filters, "status": "queued"})
|
||
|
|
running_jobs = await self.count(filters={**base_filters, "status": "running"})
|
||
|
|
completed_jobs = await self.count(filters={**base_filters, "status": "completed"})
|
||
|
|
failed_jobs = await self.count(filters={**base_filters, "status": "failed"})
|
||
|
|
cancelled_jobs = await self.count(filters={**base_filters, "status": "cancelled"})
|
||
|
|
|
||
|
|
# Get jobs by type
|
||
|
|
type_query = text(f"""
|
||
|
|
SELECT job_type, COUNT(*) as count
|
||
|
|
FROM training_job_queue
|
||
|
|
WHERE 1=1
|
||
|
|
{' AND tenant_id = :tenant_id' if tenant_id else ''}
|
||
|
|
GROUP BY job_type
|
||
|
|
ORDER BY count DESC
|
||
|
|
""")
|
||
|
|
|
||
|
|
params = {"tenant_id": tenant_id} if tenant_id else {}
|
||
|
|
result = await self.session.execute(type_query, params)
|
||
|
|
jobs_by_type = {row.job_type: row.count for row in result.fetchall()}
|
||
|
|
|
||
|
|
# Get average wait time for completed jobs
|
||
|
|
wait_time_query = text(f"""
|
||
|
|
SELECT
|
||
|
|
AVG(EXTRACT(EPOCH FROM (started_at - created_at))/60) as avg_wait_minutes
|
||
|
|
FROM training_job_queue
|
||
|
|
WHERE status = 'completed'
|
||
|
|
AND started_at IS NOT NULL
|
||
|
|
AND created_at IS NOT NULL
|
||
|
|
{' AND tenant_id = :tenant_id' if tenant_id else ''}
|
||
|
|
""")
|
||
|
|
|
||
|
|
wait_result = await self.session.execute(wait_time_query, params)
|
||
|
|
wait_row = wait_result.fetchone()
|
||
|
|
avg_wait_time = float(wait_row.avg_wait_minutes) if wait_row and wait_row.avg_wait_minutes else 0.0
|
||
|
|
|
||
|
|
return {
|
||
|
|
"tenant_id": tenant_id,
|
||
|
|
"queue_counts": {
|
||
|
|
"queued": queued_jobs,
|
||
|
|
"running": running_jobs,
|
||
|
|
"completed": completed_jobs,
|
||
|
|
"failed": failed_jobs,
|
||
|
|
"cancelled": cancelled_jobs,
|
||
|
|
"total": queued_jobs + running_jobs + completed_jobs + failed_jobs + cancelled_jobs
|
||
|
|
},
|
||
|
|
"jobs_by_type": jobs_by_type,
|
||
|
|
"avg_wait_time_minutes": round(avg_wait_time, 2),
|
||
|
|
"queue_health": {
|
||
|
|
"has_queued_jobs": queued_jobs > 0,
|
||
|
|
"has_running_jobs": running_jobs > 0,
|
||
|
|
"failure_rate": round((failed_jobs / max(completed_jobs + failed_jobs, 1)) * 100, 2)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to get queue status",
|
||
|
|
tenant_id=tenant_id,
|
||
|
|
error=str(e))
|
||
|
|
return {
|
||
|
|
"tenant_id": tenant_id,
|
||
|
|
"queue_counts": {
|
||
|
|
"queued": 0, "running": 0, "completed": 0,
|
||
|
|
"failed": 0, "cancelled": 0, "total": 0
|
||
|
|
},
|
||
|
|
"jobs_by_type": {},
|
||
|
|
"avg_wait_time_minutes": 0.0,
|
||
|
|
"queue_health": {
|
||
|
|
"has_queued_jobs": False,
|
||
|
|
"has_running_jobs": False,
|
||
|
|
"failure_rate": 0.0
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async def get_jobs_by_tenant(
|
||
|
|
self,
|
||
|
|
tenant_id: str,
|
||
|
|
status: str = None,
|
||
|
|
job_type: str = None,
|
||
|
|
skip: int = 0,
|
||
|
|
limit: int = 100
|
||
|
|
) -> List[TrainingJobQueue]:
|
||
|
|
"""Get jobs for a tenant with optional filtering"""
|
||
|
|
try:
|
||
|
|
filters = {"tenant_id": tenant_id}
|
||
|
|
if status:
|
||
|
|
filters["status"] = status
|
||
|
|
if job_type:
|
||
|
|
filters["job_type"] = job_type
|
||
|
|
|
||
|
|
return await self.get_multi(
|
||
|
|
filters=filters,
|
||
|
|
skip=skip,
|
||
|
|
limit=limit,
|
||
|
|
order_by="created_at",
|
||
|
|
order_desc=True
|
||
|
|
)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to get jobs by tenant",
|
||
|
|
tenant_id=tenant_id,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to get tenant jobs: {str(e)}")
|
||
|
|
|
||
|
|
async def cleanup_old_jobs(self, days_old: int = 30, status_filter: str = None) -> int:
|
||
|
|
"""Clean up old completed/failed/cancelled jobs"""
|
||
|
|
try:
|
||
|
|
cutoff_date = datetime.now() - timedelta(days=days_old)
|
||
|
|
|
||
|
|
# Only clean up finished jobs by default
|
||
|
|
default_statuses = ["completed", "failed", "cancelled"]
|
||
|
|
|
||
|
|
if status_filter:
|
||
|
|
status_condition = "status = :status"
|
||
|
|
params = {"cutoff_date": cutoff_date, "status": status_filter}
|
||
|
|
else:
|
||
|
|
status_list = "', '".join(default_statuses)
|
||
|
|
status_condition = f"status IN ('{status_list}')"
|
||
|
|
params = {"cutoff_date": cutoff_date}
|
||
|
|
|
||
|
|
query_text = f"""
|
||
|
|
DELETE FROM training_job_queue
|
||
|
|
WHERE created_at < :cutoff_date
|
||
|
|
AND {status_condition}
|
||
|
|
"""
|
||
|
|
|
||
|
|
result = await self.session.execute(text(query_text), params)
|
||
|
|
deleted_count = result.rowcount
|
||
|
|
|
||
|
|
logger.info("Cleaned up old queue jobs",
|
||
|
|
deleted_count=deleted_count,
|
||
|
|
days_old=days_old,
|
||
|
|
status_filter=status_filter)
|
||
|
|
|
||
|
|
return deleted_count
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to cleanup old queue jobs",
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Queue cleanup failed: {str(e)}")
|
||
|
|
|
||
|
|
async def get_stuck_jobs(self, hours_stuck: int = 2) -> List[TrainingJobQueue]:
|
||
|
|
"""Get jobs that have been running for too long"""
|
||
|
|
try:
|
||
|
|
cutoff_time = datetime.now() - timedelta(hours=hours_stuck)
|
||
|
|
|
||
|
|
query_text = """
|
||
|
|
SELECT * FROM training_job_queue
|
||
|
|
WHERE status = 'running'
|
||
|
|
AND started_at IS NOT NULL
|
||
|
|
AND started_at < :cutoff_time
|
||
|
|
ORDER BY started_at ASC
|
||
|
|
"""
|
||
|
|
|
||
|
|
result = await self.session.execute(text(query_text), {"cutoff_time": cutoff_time})
|
||
|
|
|
||
|
|
stuck_jobs = []
|
||
|
|
for row in result.fetchall():
|
||
|
|
record_dict = dict(row._mapping)
|
||
|
|
job = self.model(**record_dict)
|
||
|
|
stuck_jobs.append(job)
|
||
|
|
|
||
|
|
if stuck_jobs:
|
||
|
|
logger.warning("Found stuck jobs",
|
||
|
|
count=len(stuck_jobs),
|
||
|
|
hours_stuck=hours_stuck)
|
||
|
|
|
||
|
|
return stuck_jobs
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to get stuck jobs",
|
||
|
|
hours_stuck=hours_stuck,
|
||
|
|
error=str(e))
|
||
|
|
return []
|
||
|
|
|
||
|
|
async def reset_stuck_jobs(self, hours_stuck: int = 2) -> int:
|
||
|
|
"""Reset stuck jobs back to queued status"""
|
||
|
|
try:
|
||
|
|
stuck_jobs = await self.get_stuck_jobs(hours_stuck)
|
||
|
|
reset_count = 0
|
||
|
|
|
||
|
|
for job in stuck_jobs:
|
||
|
|
# Reset job to queued status
|
||
|
|
await self.update(job.id, {
|
||
|
|
"status": "queued",
|
||
|
|
"started_at": None,
|
||
|
|
"updated_at": datetime.now()
|
||
|
|
})
|
||
|
|
reset_count += 1
|
||
|
|
|
||
|
|
if reset_count > 0:
|
||
|
|
logger.info("Reset stuck jobs",
|
||
|
|
reset_count=reset_count,
|
||
|
|
hours_stuck=hours_stuck)
|
||
|
|
|
||
|
|
return reset_count
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Failed to reset stuck jobs",
|
||
|
|
hours_stuck=hours_stuck,
|
||
|
|
error=str(e))
|
||
|
|
raise DatabaseError(f"Failed to reset stuck jobs: {str(e)}")
|