wehub-resource-sync
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
from abc import ABC, abstractmethod
|
|
from typing import ClassVar, Protocol
|
|
|
|
|
|
class _RunnerMeta(Protocol):
|
|
INFERENCE_METHOD: ClassVar[str]
|
|
|
|
|
|
_RunnersDict = dict[str, type["_InferenceRunner"]]
|
|
|
|
|
|
# kept private until we stabilize the API (only used for EOU today)
|
|
class _InferenceRunner(ABC, _RunnerMeta):
|
|
registered_runners: _RunnersDict = {}
|
|
|
|
@classmethod
|
|
def register_runner(cls, runner_class: type[_InferenceRunner]) -> None:
|
|
if threading.current_thread() != threading.main_thread():
|
|
raise RuntimeError("InferenceRunner must be registered on the main thread")
|
|
|
|
if runner_class.INFERENCE_METHOD in cls.registered_runners:
|
|
raise ValueError(f"InferenceRunner {runner_class.INFERENCE_METHOD} already registered")
|
|
|
|
cls.registered_runners[runner_class.INFERENCE_METHOD] = runner_class
|
|
|
|
@abstractmethod
|
|
def initialize(self) -> None:
|
|
"""Initialize the runner. This is used to load models, etc."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def run(self, data: bytes) -> bytes | None:
|
|
"""Run inference on the given data."""
|
|
...
|