File size: 12,631 Bytes
1a4aa87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 | """
Priority Engine for Intelligent Job Scheduling
Computes priority scores for jobs based on multiple factors:
- Tenant plan weight (Enterprise > Pro > Basic > Free)
- SLA deadline urgency
- Cost sensitivity (inverse of budget buffer)
- Job waiting time (urgency factor)
Score = w_p * P + w_s * SLA + w_c * CostSensitivity + w_u * Urgency
"""
import uuid
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Dict, Optional
from pydantic import BaseModel
# Plan type weights as defined in Day 4 requirements
PLAN_WEIGHTS = {
"free": 0.2,
"basic": 0.4,
"pro": 0.5,
"enterprise": 1.0,
}
class PriorityQueueType(str, Enum):
"""Priority queue types for job classification."""
HIGH = "high"
MEDIUM = "medium"
LOW = "low"
class JobPriorityScore(BaseModel):
"""Complete priority score breakdown for a job."""
job_id: uuid.UUID
tenant_id: uuid.UUID
# Component scores (all normalized to [0, 1])
plan_score: float = 0.0 # Based on tenant plan
sla_score: float = 0.0 # Based on deadline urgency
cost_sensitivity_score: float = 0.0 # Based on budget buffer
urgency_score: float = 0.0 # Based on wait time
# Weighted total
total_score: float = 0.0
# Assigned queue
queue_type: PriorityQueueType = PriorityQueueType.MEDIUM
# Metadata
calculated_at: datetime = None
estimated_gpu_hours: float = 0.0
estimated_cost: float = 0.0
def __init__(self, **data):
if "calculated_at" not in data or data["calculated_at"] is None:
data["calculated_at"] = datetime.utcnow()
super().__init__(**data)
class PriorityEngine:
"""
Multi-factor priority scoring engine for job scheduling.
Computes priority scores based on:
- Tenant plan tier (Enterprise gets scheduling preference)
- SLA deadline urgency (closer deadlines = higher priority)
- Cost sensitivity (lower budget buffer = higher priority)
- Job waiting time (longer waits = higher priority)
"""
# Default weights for priority components
DEFAULT_WEIGHTS = {
"plan": 0.30, # w_p
"sla": 0.25, # w_s
"cost_sensitivity": 0.20, # w_c
"urgency": 0.25, # w_u
}
# Queue score thresholds
HIGH_QUEUE_THRESHOLD = 0.7
MEDIUM_QUEUE_THRESHOLD = 0.4
def __init__(
self,
weights: Optional[Dict[str, float]] = None,
avg_inference_time_ms: float = 500.0,
):
"""
Initialize priority engine.
Args:
weights: Custom weights for priority components.
If None, uses DEFAULT_WEIGHTS.
avg_inference_time_ms: Average inference time in milliseconds.
Used for GPU hour estimation.
"""
self.weights = weights or self.DEFAULT_WEIGHTS
self.avg_inference_time_ms = avg_inference_time_ms
# Ensure weights sum to 1.0
total = sum(self.weights.values())
if total != 1.0:
# Normalize weights
self.weights = {k: v / total for k, v in self.weights.items()}
def calculate_priority_score(
self,
job_id: uuid.UUID,
tenant_id: uuid.UUID,
plan_type: str,
total_samples: int,
deadline_timestamp: Optional[datetime] = None,
budget_limit: Optional[float] = None,
used_budget: Optional[float] = None,
submitted_at: Optional[datetime] = None,
current_time: Optional[datetime] = None,
) -> JobPriorityScore:
"""
Calculate comprehensive priority score for a job.
Args:
job_id: Unique job identifier
tenant_id: Tenant identifier
plan_type: Tenant plan type (free, basic, pro, enterprise)
total_samples: Number of samples to process
deadline_timestamp: Optional deadline for SLA enforcement
budget_limit: Optional budget limit for the tenant
used_budget: Amount of budget already used
submitted_at: When the job was submitted
current_time: Current timestamp (defaults to now)
Returns:
JobPriorityScore with all component scores and total
"""
now = current_time or datetime.utcnow()
# 1. Calculate plan score (based on tenant tier)
plan_score = self._calculate_plan_score(plan_type)
# 2. Calculate SLA score (based on deadline urgency)
sla_score = self._calculate_sla_score(deadline_timestamp, now)
# 3. Calculate cost sensitivity (based on budget buffer)
cost_sensitivity = self._calculate_cost_sensitivity(
budget_limit, used_budget
)
# 4. Calculate urgency score (based on wait time)
urgency_score = self._calculate_urgency_score(submitted_at, now)
# Calculate weighted total
total_score = (
self.weights["plan"] * plan_score +
self.weights["sla"] * sla_score +
self.weights["cost_sensitivity"] * cost_sensitivity +
self.weights["urgency"] * urgency_score
)
# Determine queue type
queue_type = self._determine_queue_type(total_score)
# Estimate GPU hours and cost
estimated_gpu_hours = self._estimate_gpu_hours(total_samples)
estimated_cost = self._estimate_job_cost(estimated_gpu_hours)
return JobPriorityScore(
job_id=job_id,
tenant_id=tenant_id,
plan_score=plan_score,
sla_score=sla_score,
cost_sensitivity_score=cost_sensitivity,
urgency_score=urgency_score,
total_score=total_score,
queue_type=queue_type,
calculated_at=now,
estimated_gpu_hours=estimated_gpu_hours,
estimated_cost=estimated_cost,
)
def _calculate_plan_score(self, plan_type: str) -> float:
"""
Calculate score based on tenant plan type.
Returns normalized score based on PLAN_WEIGHTS.
"""
return PLAN_WEIGHTS.get(plan_type.lower(), PLAN_WEIGHTS["free"])
def _calculate_sla_score(
self,
deadline_timestamp: Optional[datetime],
current_time: datetime,
) -> float:
"""
Calculate SLA urgency score.
Jobs with imminent deadlines get higher priority.
If no deadline, returns 0.5 (medium priority).
"""
if deadline_timestamp is None:
return 0.5 # No deadline = medium priority
time_remaining = deadline_timestamp - current_time
if time_remaining.total_seconds() <= 0:
# Deadline passed - maximum urgency
return 1.0
# Score decreases as more time remains
# Map to [0, 1] where 1 is urgent (very little time)
max_urgency_hours = 1.0 # 1 hour = maximum urgency
min_urgency_hours = 24.0 # 24 hours = minimum urgency
hours_remaining = time_remaining.total_seconds() / 3600
if hours_remaining <= max_urgency_hours:
return 1.0
elif hours_remaining >= min_urgency_hours:
return 0.0
else:
# Linear interpolation
return 1.0 - (hours_remaining - max_urgency_hours) / (min_urgency_hours - max_urgency_hours)
def _calculate_cost_sensitivity(
self,
budget_limit: Optional[float],
used_budget: Optional[float],
) -> float:
"""
Calculate cost sensitivity based on budget buffer.
Returns higher score when budget is nearly exhausted.
"""
if budget_limit is None or used_budget is None:
return 0.5 # No budget tracking = medium sensitivity
if budget_limit <= 0:
return 1.0 # No budget = maximum sensitivity
# Calculate buffer remaining
buffer_remaining = budget_limit - used_budget
buffer_percent = buffer_remaining / budget_limit
if buffer_percent <= 0:
return 1.0 # Budget exhausted = maximum sensitivity
elif buffer_percent >= 1.0:
return 0.0 # Full budget = no sensitivity
else:
# Higher sensitivity as buffer decreases
return 1.0 - buffer_percent
def _calculate_urgency_score(
self,
submitted_at: Optional[datetime],
current_time: datetime,
) -> float:
"""
Calculate urgency based on job waiting time.
Jobs waiting longer get higher priority to prevent starvation.
"""
if submitted_at is None:
return 0.5 # No submission time = medium urgency
wait_time = current_time - submitted_at
wait_hours = wait_time.total_seconds() / 3600
# Maximum urgency after 4 hours of waiting
# Minimum urgency (0) for jobs submitted < 5 minutes ago
if wait_hours >= 4.0:
return 1.0
elif wait_hours <= 0.083: # 5 minutes
return 0.0
else:
# Linear interpolation
return (wait_hours - 0.083) / (4.0 - 0.083)
def _determine_queue_type(self, total_score: float) -> PriorityQueueType:
"""Determine which priority queue a job belongs to."""
if total_score >= self.HIGH_QUEUE_THRESHOLD:
return PriorityQueueType.HIGH
elif total_score >= self.MEDIUM_QUEUE_THRESHOLD:
return PriorityQueueType.MEDIUM
else:
return PriorityQueueType.LOW
def _estimate_gpu_hours(self, total_samples: int) -> float:
"""
Estimate GPU hours required for a job.
GPUHours = (Samples * AvgInferenceTime) / 3600
"""
# Average time per sample in seconds
avg_time_per_sample_seconds = self.avg_inference_time_ms / 1000.0
# Total time in seconds
total_time_seconds = total_samples * avg_time_per_sample_seconds
# Convert to hours
gpu_hours = total_time_seconds / 3600.0
return gpu_hours
def _estimate_job_cost(self, estimated_gpu_hours: float) -> float:
"""
Estimate cost based on GPU hours.
Cost = GPUHours * CostPerGPUHour
Default: $3.00 per GPU hour (AWS p3.2xlarge spot price)
"""
cost_per_gpu_hour = 3.00 # Default cost
return estimated_gpu_hours * cost_per_gpu_hour
def should_promote_priority(
self,
current_score: JobPriorityScore,
current_time: Optional[datetime] = None,
) -> bool:
"""
Check if a job should be promoted to higher priority.
Promotes if:
- Deadline is approaching and job is not in high queue
- Job has been waiting too long (starvation prevention)
"""
now = current_time or datetime.utcnow()
# Calculate time since submission
if current_score.calculated_at:
wait_time = now - current_score.calculated_at
wait_hours = wait_time.total_seconds() / 3600
# Promote if waiting > 2 hours in low queue
if (current_score.queue_type == PriorityQueueType.LOW and
wait_hours > 2.0):
return True
# Promote if SLA is at risk
if (current_score.sla_score > 0.8 and
current_score.queue_type != PriorityQueueType.HIGH):
return True
return False
# Global instance
_priority_engine: Optional[PriorityEngine] = None
def get_priority_engine() -> PriorityEngine:
"""Get or create the global PriorityEngine instance."""
global _priority_engine
if _priority_engine is None:
_priority_engine = PriorityEngine()
return _priority_engine
__all__ = [
"PriorityEngine",
"JobPriorityScore",
"PriorityQueueType",
"PLAN_WEIGHTS",
"get_priority_engine",
]
|