File size: 2,264 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 | """
Queue Module for AegisLM
Provides asynchronous job queue functionality for enterprise-grade
evaluation processing with checkpointing and progress tracking.
Components:
- job_schema: Job definitions and schemas
- worker_schema: Worker definitions and schemas
- status_tracker: Job status tracking and state management
- producer: Job submission producers
- consumer: Worker that processes jobs from the queue
- worker_registry: Worker registration and management
- scheduler: GPU-aware job scheduling
"""
from .job_schema import (
JobStatus,
JobType,
JobPriority,
GPURequirement,
EvaluationJob,
JobSubmissionRequest,
JobStatusResponse,
JobProgressUpdate,
)
from .worker_schema import (
WorkerStatus,
GPUInfo,
WorkerRegistrationRequest,
WorkerRegistrationResponse,
HeartbeatRequest,
HeartbeatResponse,
WorkerStatusResponse,
WorkerListResponse,
WorkerMetricsResponse,
)
from .status_tracker import (
JobStatusTracker,
status_tracker,
get_status_tracker,
)
from .producer import (
JobProducer,
get_job_producer,
)
from .consumer import (
EvaluationWorker,
WorkerConfig,
get_worker,
)
from .worker_registry import (
WorkerRegistry,
get_worker_registry,
)
from .scheduler import (
JobScheduler,
get_job_scheduler,
)
__all__ = [
# Job schema
"JobStatus",
"JobType",
"JobPriority",
"GPURequirement",
"EvaluationJob",
"JobSubmissionRequest",
"JobStatusResponse",
"JobProgressUpdate",
# Worker schema
"WorkerStatus",
"GPUInfo",
"WorkerRegistrationRequest",
"WorkerRegistrationResponse",
"HeartbeatRequest",
"HeartbeatResponse",
"WorkerStatusResponse",
"WorkerListResponse",
"WorkerMetricsResponse",
# Status tracker
"JobStatusTracker",
"status_tracker",
"get_status_tracker",
# Producer
"JobProducer",
"get_job_producer",
# Consumer
"EvaluationWorker",
"WorkerConfig",
"get_worker",
# Worker registry
"WorkerRegistry",
"get_worker_registry",
# Scheduler
"JobScheduler",
"get_job_scheduler",
]
|