50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
from dataclasses import dataclass, field
|
|
import asyncio
|
|
import time
|
|
|
|
|
|
class BaseJobState:
|
|
"""Marker base class for typed USDU job state containers."""
|
|
|
|
|
|
@dataclass
|
|
class TileJobState(BaseJobState):
|
|
"""Typed state container for static (tile) USDU jobs."""
|
|
|
|
multi_job_id: str
|
|
mode: str = field(default="static", init=False)
|
|
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
completed_tasks: dict = field(default_factory=dict)
|
|
worker_status: dict = field(default_factory=dict)
|
|
assigned_to_workers: dict = field(default_factory=dict)
|
|
batch_size: int = 0
|
|
num_tiles_per_image: int = 0
|
|
batched_static: bool = False
|
|
created_at: float = field(default_factory=time.monotonic)
|
|
|
|
|
|
@dataclass
|
|
class ImageJobState(BaseJobState):
|
|
"""Typed state container for dynamic (per-image) USDU jobs."""
|
|
|
|
multi_job_id: str
|
|
mode: str = field(default="dynamic", init=False)
|
|
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
completed_images: dict = field(default_factory=dict)
|
|
worker_status: dict = field(default_factory=dict)
|
|
assigned_to_workers: dict = field(default_factory=dict)
|
|
batch_size: int = 0
|
|
num_tiles_per_image: int = 0
|
|
batched_static: bool = False
|
|
created_at: float = field(default_factory=time.monotonic)
|
|
|
|
@property
|
|
def pending_tasks(self):
|
|
return self.pending_images
|
|
|
|
@property
|
|
def completed_tasks(self):
|
|
return self.completed_images
|