backend/queue/__init__.py

"""
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",
]