src/errors.py

from __future__ import annotations

from dataclasses import dataclass


INPUT_VALIDATION = "INPUT_VALIDATION"
MEDIA_PREPARATION = "MEDIA_PREPARATION"
MODEL_DOWNLOAD = "MODEL_DOWNLOAD"
MODEL_LOAD = "MODEL_LOAD"
GPU_INFERENCE = "GPU_INFERENCE"
OUTPUT_ENCODING = "OUTPUT_ENCODING"
BUNDLE_PACKAGING = "BUNDLE_PACKAGING"
PREPARED_JOB_VALIDATION = "PREPARED_JOB_VALIDATION"


@dataclass
class StageError(RuntimeError):
    stage: str
    message: str
    retryable: bool = True
    hint: str = ""

    def __post_init__(self) -> None:
        RuntimeError.__init__(self, self.message)

    def to_dict(self) -> dict:
        return {
            "stage": self.stage,
            "error_type": type(self).__name__,
            "message": self.message,
            "retryable": bool(self.retryable),
            "hint": self.hint or None,
        }


def stage_error(stage: str, exc: Exception, *, hint: str = "", retryable: bool = True) -> StageError:
    if isinstance(exc, StageError):
        return exc
    return StageError(stage=stage, message=str(exc), hint=hint, retryable=retryable)