Spaces:
Running on Zero
Running on Zero
File size: 5,931 Bytes
5f7afef | 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 | from __future__ import annotations
import os
import time
import uuid
import threading
from concurrent.futures import ThreadPoolExecutor
from enum import Enum
from typing import Dict, Optional, Any
from pydantic import BaseModel
from codestar.core.schemas import JudgeResult, JudgeStatus
from codestar.core.engine import JudgeEngine
from codestar.problems.loader import repo
from codestar.server.db import create_submission, calculate_beats
class TaskStatus(str, Enum):
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
class JudgeTaskInfo(BaseModel):
task_id: str
problem_slug: str
status: TaskStatus
created_at: float
started_at: Optional[float] = None
completed_at: Optional[float] = None
result: Optional[JudgeResult] = None
error: Optional[str] = None
class JudgeTaskQueue:
def __init__(self, max_workers: Optional[int] = None, max_tasks: int = 500):
workers = max_workers or int(os.environ.get("JUDGE_WORKERS", 4))
self._executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="judge-worker")
self._tasks: Dict[str, JudgeTaskInfo] = {}
self._max_tasks = max_tasks
self._lock = threading.Lock()
self._engine = JudgeEngine()
def submit_task(
self,
problem_slug: str,
code: str,
is_test: bool = False,
custom_testcase: Optional[str] = None,
user_id: Optional[int] = None,
) -> str:
task_id = f"task_{uuid.uuid4().hex[:12]}"
now = time.time()
task_info = JudgeTaskInfo(
task_id=task_id,
problem_slug=problem_slug,
status=TaskStatus.PENDING,
created_at=now,
)
with self._lock:
self._cleanup_old_tasks()
self._tasks[task_id] = task_info
self._executor.submit(
self._execute_task,
task_id=task_id,
problem_slug=problem_slug,
code=code,
is_test=is_test,
custom_testcase=custom_testcase,
user_id=user_id,
)
return task_id
def _execute_task(
self,
task_id: str,
problem_slug: str,
code: str,
is_test: bool,
custom_testcase: Optional[str],
user_id: Optional[int],
):
with self._lock:
if task_id not in self._tasks:
return
self._tasks[task_id].status = TaskStatus.RUNNING
self._tasks[task_id].started_at = time.time()
try:
problem = repo.get(problem_slug)
if not problem:
with self._lock:
self._tasks[task_id].status = TaskStatus.FAILED
self._tasks[task_id].error = f"Problem '{problem_slug}' not found"
self._tasks[task_id].completed_at = time.time()
return
result = self._engine.judge(problem, code, custom_testcase=custom_testcase)
# Calculate beats for AC submissions
if result.status == JudgeStatus.AC:
beats = calculate_beats(problem_slug, result.execution_time_ms, result.peak_memory_mb)
result.beats_runtime_pct = beats["runtime_pct"]
result.beats_memory_pct = beats["memory_pct"]
# Persist submission if not test
if not is_test:
passed_cases = sum(1 for c in result.cases if c.passed)
total_cases = len(result.cases)
status_str = result.status.value if hasattr(result.status, "value") else str(result.status)
err_msg = result.message if result.status != JudgeStatus.AC else None
try:
create_submission(
user_id=user_id,
problem_slug=problem_slug,
status=status_str,
passed_cases=passed_cases,
total_cases=total_cases,
runtime_ms=result.execution_time_ms,
code=code,
error_message=err_msg,
peak_memory_mb=result.peak_memory_mb,
)
except Exception:
pass
with self._lock:
self._tasks[task_id].status = TaskStatus.COMPLETED
self._tasks[task_id].result = result
self._tasks[task_id].completed_at = time.time()
except Exception as e:
with self._lock:
self._tasks[task_id].status = TaskStatus.FAILED
self._tasks[task_id].error = str(e)
self._tasks[task_id].completed_at = time.time()
def get_task(self, task_id: str) -> Optional[JudgeTaskInfo]:
with self._lock:
return self._tasks.get(task_id)
def _cleanup_old_tasks(self, max_age_seconds: float = 3600.0):
"""Clean up tasks completed more than an hour ago, or evict oldest completed tasks if exceeding max capacity."""
now = time.time()
# 1. Clean up expired tasks
to_delete = [
tid for tid, info in self._tasks.items()
if info.completed_at and (now - info.completed_at > max_age_seconds)
]
for tid in to_delete:
del self._tasks[tid]
# 2. Strict capacity bounding: if still exceeding max_tasks, evict oldest completed
if len(self._tasks) > self._max_tasks:
completed_tasks = sorted(
[info for info in self._tasks.values() if info.completed_at],
key=lambda x: x.completed_at or 0.0,
)
excess = len(self._tasks) - self._max_tasks
for info in completed_tasks[:excess]:
self._tasks.pop(info.task_id, None)
# Global singleton instance
judge_queue = JudgeTaskQueue()
|