46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
import importlib.util
|
|
import subprocess
|
|
import sys
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from source.manager import ModelManager
|
|
from source.model import base
|
|
|
|
|
|
class PythonModel(base.BaseModel):
|
|
"""
|
|
A model running a custom python model.
|
|
"""
|
|
|
|
def __init__(self, manager: ModelManager, configuration: dict, path: Path):
|
|
super().__init__(manager, configuration, path)
|
|
|
|
# get the name of the file containing the model code
|
|
file = configuration.get("file")
|
|
if file is None:
|
|
raise ValueError("Field 'file' is missing from the configuration")
|
|
|
|
# install custom requirements
|
|
requirements = configuration.get("requirements", [])
|
|
if len(requirements) > 0:
|
|
subprocess.run([sys.executable, "-m", "pip", "install", *requirements])
|
|
|
|
# create the module specification
|
|
module_spec = importlib.util.spec_from_file_location(
|
|
f"model-{uuid.uuid4()}",
|
|
self.path / file
|
|
)
|
|
# get the module
|
|
self.module = importlib.util.module_from_spec(module_spec)
|
|
# load the module
|
|
module_spec.loader.exec_module(self.module)
|
|
|
|
def _load(self) -> None:
|
|
return self.module.load(self)
|
|
|
|
def _unload(self) -> None:
|
|
return self.module.unload(self)
|
|
|
|
def _infer(self, payload: dict) -> str | bytes:
|
|
return self.module.infer(self, payload)
|