2023-02-26 05:49:39 +00:00
|
|
|
from logging import getLogger
|
2023-02-26 05:55:30 +00:00
|
|
|
from typing import Any, Callable, Tuple
|
2023-02-26 05:49:39 +00:00
|
|
|
|
2023-02-26 20:15:30 +00:00
|
|
|
from torch.multiprocessing import Queue, Value
|
|
|
|
|
2023-02-26 05:49:39 +00:00
|
|
|
from ..params import DeviceParams
|
|
|
|
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
ProgressCallback = Callable[[int, int, Any], None]
|
|
|
|
|
2023-02-26 20:15:30 +00:00
|
|
|
|
2023-02-26 05:49:39 +00:00
|
|
|
class WorkerContext:
|
|
|
|
cancel: "Value[bool]" = None
|
2023-02-27 23:35:31 +00:00
|
|
|
job: str = None
|
2023-02-26 05:55:30 +00:00
|
|
|
pending: "Queue[Tuple[Callable, Any, Any]]" = None
|
2023-02-26 05:49:39 +00:00
|
|
|
progress: "Value[int]" = None
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
self,
|
2023-02-27 23:35:31 +00:00
|
|
|
job: str,
|
2023-02-26 05:49:39 +00:00
|
|
|
device: DeviceParams,
|
2023-02-26 21:06:40 +00:00
|
|
|
cancel: "Value[bool]" = None,
|
|
|
|
logs: "Queue[str]" = None,
|
|
|
|
pending: "Queue[Any]" = None,
|
2023-02-27 02:37:22 +00:00
|
|
|
progress: "Queue[Tuple[str, int]]" = None,
|
|
|
|
finished: "Queue[str]" = None,
|
2023-02-26 05:49:39 +00:00
|
|
|
):
|
2023-02-27 23:35:31 +00:00
|
|
|
self.job = job
|
2023-02-26 05:49:39 +00:00
|
|
|
self.device = device
|
2023-02-27 02:09:42 +00:00
|
|
|
self.cancel = cancel
|
2023-02-26 05:49:39 +00:00
|
|
|
self.progress = progress
|
2023-02-26 18:51:11 +00:00
|
|
|
self.finished = finished
|
2023-02-27 02:09:42 +00:00
|
|
|
self.logs = logs
|
|
|
|
self.pending = pending
|
2023-02-26 05:49:39 +00:00
|
|
|
|
|
|
|
def is_cancelled(self) -> bool:
|
|
|
|
return self.cancel.value
|
|
|
|
|
|
|
|
def get_device(self) -> DeviceParams:
|
|
|
|
"""
|
|
|
|
Get the device assigned to this job.
|
|
|
|
"""
|
|
|
|
return self.device
|
|
|
|
|
|
|
|
def get_progress(self) -> int:
|
|
|
|
return self.progress.value
|
|
|
|
|
|
|
|
def get_progress_callback(self) -> ProgressCallback:
|
|
|
|
def on_progress(step: int, timestep: int, latents: Any):
|
|
|
|
on_progress.step = step
|
|
|
|
if self.is_cancelled():
|
|
|
|
raise RuntimeError("job has been cancelled")
|
|
|
|
else:
|
2023-02-27 23:35:31 +00:00
|
|
|
logger.debug("setting progress for job %s to %s", self.job, step)
|
2023-02-26 05:49:39 +00:00
|
|
|
self.set_progress(step)
|
|
|
|
|
|
|
|
return on_progress
|
|
|
|
|
|
|
|
def set_cancel(self, cancel: bool = True) -> None:
|
|
|
|
with self.cancel.get_lock():
|
|
|
|
self.cancel.value = cancel
|
|
|
|
|
|
|
|
def set_progress(self, progress: int) -> None:
|
2023-02-28 04:37:43 +00:00
|
|
|
self.progress.put((self.job, self.device.device, progress), block=False)
|
2023-02-27 02:09:42 +00:00
|
|
|
|
2023-02-27 02:37:22 +00:00
|
|
|
def set_finished(self) -> None:
|
2023-02-27 23:35:31 +00:00
|
|
|
self.finished.put((self.job, self.device.device))
|
2023-02-27 02:09:42 +00:00
|
|
|
|
2023-02-26 21:06:40 +00:00
|
|
|
def clear_flags(self) -> None:
|
|
|
|
self.set_cancel(False)
|
|
|
|
self.set_progress(0)
|
2023-02-27 23:14:53 +00:00
|
|
|
|
|
|
|
|
|
|
|
class JobStatus:
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
name: str,
|
|
|
|
progress: int = 0,
|
|
|
|
cancelled: bool = False,
|
|
|
|
finished: bool = False,
|
|
|
|
) -> None:
|
|
|
|
self.name = name
|
|
|
|
self.progress = progress
|
|
|
|
self.cancelled = cancelled
|
|
|
|
self.finished = finished
|