Spaces:
Running on Zero
Running on Zero
| 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() | |