fnnx 0.0.1__py3-none-any.whl

fnnx/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

fnnx/device.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class DeviceConfig:
+    accelerator: str
+    device_config: dict | None
+
+
+@dataclass
+class DeviceMap:
+    accelerator: str
+    node_device_map: dict[str, dict]
+    variant_device_config: dict | str | None = None

fnnx/dtypes.py

@@ -0,0 +1,191 @@
+from copy import deepcopy
+import jsonschema
+from copy import copy, deepcopy
+
+
+class FlatList:
+    def __init__(self, data: list):
+        if not isinstance(data, list):
+            raise ValueError("FlatList only accepts lists")
+        self.data = copy(data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def __setitem__(self, index, value):
+        self.data[index] = value
+
+    def __repr__(self):
+        return f"FlatList({self.data})"
+
+    def append(self, value):
+        self.data.append(value)
+
+
+class DtypesManager:
+
+    def __init__(self, external_dtypes: dict, builtins: dict):
+        self.dtypes = deepcopy(external_dtypes)
+        self.dtypes.update(deepcopy(builtins))
+        for dtype in self.dtypes:
+            if "[" in dtype:
+                raise ValueError(f"Invalid dtype name: {dtype}")
+        for reserved_type in [
+            "string",
+            "integer",
+            "float",
+            "Array",
+            "NDContainer",
+            "FlatList",
+        ]:
+            if reserved_type in self.dtypes:
+                raise ValueError(f"Invalid dtype name: {reserved_type}")
+
+    def get_dtype(self, dtype_name: str):
+        if dtype_name not in self.dtypes:
+            raise ValueError(f"Unknown dtype: {dtype_name}")
+        return self.dtypes[dtype_name]
+
+    def _validate_dtype(self, dtype_name: str, data: dict):
+        schema = self.get_dtype(dtype_name)
+        jsonschema.validate(data, schema)
+
+    def validate_dtype(self, dtype_name: str, data):
+        if isinstance(data, list):
+            for d in data:
+                self.validate_dtype(dtype_name, d)
+        elif isinstance(data, FlatList):
+            self.validate_dtype(dtype_name, data.data)
+        elif isinstance(data, dict):
+            self._validate_dtype(dtype_name, data)
+        elif isinstance(data, str):
+            if dtype_name != "string":
+                raise TypeError(
+                    f"Invalid data type, expected `string`, got `{dtype_name}`"
+                )
+        elif isinstance(data, int):
+            if dtype_name != "integer":
+                raise TypeError(
+                    f"Invalid data type, expected `integer`, got `{dtype_name}`"
+                )
+        elif isinstance(data, float):
+            if dtype_name != "float":
+                raise TypeError(
+                    f"Invalid data type, expected `float`, got `{dtype_name}`"
+                )
+        else:
+            raise TypeError(f"Invalid data type: {type(data)}")
+
+
+class NDContainer:
+    def __init__(self, data, dtype, dtypes_manager: DtypesManager):
+
+        if dtype.startswith("Array["):
+            raise ValueError("NDContainer does not support Array dtype")
+        elif dtype.startswith("NDContainer["):
+            dtype = dtype[12:-1]
+
+        self.data = deepcopy(data if isinstance(data, list) else [data])
+
+        if dtypes_manager:
+            dtypes_manager.validate_dtype(dtype, self.data)
+        self.dtypes_manager = dtypes_manager
+        self._dtype = dtype
+        if "FlatList[" in self._dtype:
+            self.data = self._inner_to_flatlist(self.data)
+        self.shape = tuple(self._compute_shape(self.data))
+
+    def _inner_to_flatlist(self, nested_list, root=True):
+        if root:
+            # check base case
+            if all(not isinstance(item, (list, FlatList)) for item in nested_list):
+                return [FlatList(nested_list)]
+
+        for i, item in enumerate(nested_list):
+            if isinstance(item, list):
+                if all(not isinstance(sub_item, (list, FlatList)) for sub_item in item):
+                    nested_list[i] = FlatList(item)
+                else:
+                    self._inner_to_flatlist(item, root=False)
+        return nested_list
+
+    def _compute_shape(self, data):
+        if not isinstance(data, list) or not data:
+            return []
+        sub_shape = self._compute_shape(data[0])
+        return [len(data)] + sub_shape
+
+    def __getitem__(self, index):
+        if isinstance(index, tuple):
+            result = self.data
+            for idx in index:
+                result = result[idx]
+            return result
+        return self.data[index]
+
+    def reshape(self, *new_shape):
+        if isinstance(new_shape[0], tuple) or isinstance(new_shape[0], list):
+            new_shape = new_shape[0]
+        # Check if the total number of elements matches
+        if self._product(new_shape) != self._product(self.shape):
+            raise ValueError(
+                "Cannot reshape array of size {} into shape {}".format(
+                    self._product(self.shape), new_shape
+                )
+            )
+        flat_list = self.flatten(self.data)
+        if self._product(new_shape) != self._product(self.shape):
+            raise ValueError(
+                "Cannot reshape array of size {} into shape {}".format(
+                    self._product(self.shape), new_shape
+                )
+            )
+        reshaped = self._reshape_helper(flat_list, list(new_shape))
+        return NDContainer(
+            data=reshaped, dtype=self.dtype, dtypes_manager=self.dtypes_manager
+        )
+
+    def flatten(self, data: list | None = None):
+        return NDContainer(
+            data=self._flatten(data or self.data),
+            dtype=self.dtype,
+            dtypes_manager=self.dtypes_manager,
+        )
+
+    def _flatten(self, data) -> list:
+        result = []
+        for item in data:
+            if isinstance(item, list):
+                result.extend(self._flatten(item))
+            else:
+                result.append(item)
+        return result
+
+    def _reshape_helper(self, flat_list, shape):
+        if len(shape) == 1:
+            return flat_list[: shape[0]]
+        step = self._product(shape[1:])
+        return [
+            self._reshape_helper(flat_list[i * step : (i + 1) * step], shape[1:])
+            for i in range(shape[0])
+        ]
+
+    def _product(self, shape):
+        product = 1
+        for dim in shape:
+            product *= dim
+        return product
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @dtype.setter
+    def dtype(self, value):
+        raise AttributeError("Cannot modify immutable attribute dtype")
+
+    def __repr__(self) -> str:
+        return f"NDContainer(shape={self.shape}, dtype={self._dtype}, data={self.data})"
+
+
+BUILTINS = {}

fnnx/handlers/_base.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from fnnx.device import DeviceMap
+from dataclasses import dataclass
+
+
+@dataclass
+class BaseHandlerConfig:
+    pass
+
+
+class BaseHandler(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        model_path: str,
+        device_map: DeviceMap,
+        handler_config: BaseHandlerConfig | None = None,
+        **kwargs,
+    ):
+        pass
+
+    @abstractmethod
+    def compute(self, inputs: dict, dynamic_attributes: dict) -> dict:
+        pass
+
+    @abstractmethod
+    async def compute_async(self, inputs: dict, dynamic_attributes: dict) -> dict:
+        pass

fnnx/handlers/_common.py

@@ -0,0 +1,12 @@
+import os
+import tarfile
+import tempfile
+
+
+def unpack_model(model_path: str) -> tuple[str, bool]:
+    if os.path.isdir(model_path):
+        return model_path, False
+    with tarfile.open(model_path, "r") as tar:
+        tmp_dir = tempfile.mkdtemp(prefix="fnnx_")  #
+        tar.extractall(tmp_dir, filter="data")
+    return tmp_dir, True

fnnx/handlers/local.py

@@ -0,0 +1,167 @@
+try:
+    import numpy as np
+except ImportError:
+    np = None
+from os.path import join as pjoin
+from shutil import rmtree
+import atexit
+import json
+from dataclasses import dataclass
+from concurrent.futures import ThreadPoolExecutor
+from typing import Type
+from fnnx.device import DeviceMap
+from fnnx.dtypes import DtypesManager, BUILTINS, NDContainer
+from fnnx.variants.pipeline import Pipeline
+from fnnx.handlers._base import BaseHandler, BaseHandlerConfig
+from fnnx.handlers._common import unpack_model
+from fnnx.registry import Registry
+from fnnx.variants.pyfunc import PyFuncVariant
+from fnnx.validators.model_schema import (
+    validate_manifest,
+    validate_op_instances,
+    validate_variant,
+)
+from fnnx.ops._base import BaseOp
+from fnnx.variants._base import BaseVariant
+
+
+@dataclass
+class LocalHandlerConfig(BaseHandlerConfig):
+    n_workers: int = 1
+    n_workers_node: int = 1
+    auto_cleanup: bool = True
+    extra_ops: dict[str, Type[BaseOp]] | None = None
+
+
+class LocalHandler(BaseHandler):
+
+    def __init__(
+        self,
+        model_path: str,
+        device_map: DeviceMap,
+        handler_config: LocalHandlerConfig | None = None,
+        **kwargs,
+    ):
+        if handler_config is None:
+            handler_config = LocalHandlerConfig()
+
+        if not isinstance(device_map, DeviceMap):
+            raise ValueError("device_map must be an instance of DeviceMap")
+
+        model_path, cleanup = unpack_model(model_path)
+
+        self.model_path = model_path
+        self.cleanup = handler_config.auto_cleanup and cleanup
+
+        # should this be done on exit or on delete?
+        if self.cleanup:
+            # passing model_path and not self.model_path to avoid reference on self
+            atexit.register(lambda: _cleanup(model_path))
+
+        with open(pjoin(self.model_path, "manifest.json"), "r") as f:
+            self.manifest = json.load(f)
+            validate_manifest(self.manifest)
+
+            self.input_specs = {spec["name"]: spec for spec in self.manifest["inputs"]}
+            self.output_specs = {
+                spec["name"]: spec for spec in self.manifest["outputs"]
+            }
+
+        with open(pjoin(self.model_path, "ops.json"), "r") as f:
+            self.ops = json.load(f)
+            validate_op_instances(self.ops)
+
+        with open(pjoin(self.model_path, "variant_config.json"), "r") as f:
+            self.variant_config = json.load(f)
+            validate_variant(self.manifest["variant"], self.variant_config)
+
+        with open(pjoin(self.model_path, "dtypes.json"), "r") as f:
+            external_dtypes = json.load(f)
+            self.dtypes_manager = DtypesManager(external_dtypes, BUILTINS)
+
+        variant = self.manifest.get("variant")
+
+        registry = Registry()
+        registry.register_default_ops()
+        if handler_config.extra_ops:
+            for op_name, op in handler_config.extra_ops.items():
+                registry.register_op(op, op_name)
+
+        if variant == "pipeline":
+            vcls = Pipeline
+        elif variant == "pyfunc":
+            vcls = PyFuncVariant
+        else:
+            raise ValueError(f"Unknown variant: {variant}")
+
+        self.executor = ThreadPoolExecutor(max_workers=handler_config.n_workers)
+        self.op_executor = ThreadPoolExecutor(max_workers=handler_config.n_workers_node)
+        self.vrt: BaseVariant = vcls(
+            self.model_path,
+            self.ops,
+            self.variant_config,
+            registry=registry,
+            device_map=device_map,
+            dtypes_manager=self.dtypes_manager,
+            executor=self.executor,
+            op_executor=self.op_executor,
+        ).warmup()
+
+    def _prepare_inputs(self, inputs):
+        prepared_inputs = {}
+        for name, input in inputs.items():
+            spec = self.input_specs[name]
+            if spec["content_type"] == "NDJSON":
+                if "NDContainer[" in spec["dtype"]:
+                    if not isinstance(input, NDContainer):
+                        prepared_inputs[name] = NDContainer(
+                            input,
+                            dtype=spec["dtype"],
+                            dtypes_manager=self.dtypes_manager,
+                        )
+                    else:
+                        prepared_inputs[name] = input
+                elif "Array[" in spec["dtype"]:
+                    if np is None:
+                        raise RuntimeError(
+                            "You must have numpy installed to use Array dtype"
+                        )
+                    dtype = spec["dtype"][6:-1]
+                    if dtype == "string":
+                        prepared_inputs[name] = np.asarray(input).astype(np.str_)
+                    else:
+                        prepared_inputs[name] = np.asarray(input).astype(dtype)
+                else:
+                    raise ValueError(f"Unknown dtype {spec['dtype']}")
+            else:
+                raise ValueError(f"Unknown input type {spec['content_type']}")
+        return prepared_inputs
+
+    def _prepare_outputs(self, outputs: dict) -> dict:
+        return {k: outputs[k] for k in self.output_specs.keys()}
+
+    def compute(self, inputs: dict, dynamic_attributes: dict) -> dict:
+        res = self.vrt.compute(
+            self._prepare_inputs(inputs),
+            dynamic_attributes=dynamic_attributes,
+        )
+        return self._prepare_outputs(res)
+
+    async def compute_async(self, inputs: dict, dynamic_attributes: dict) -> dict:
+        res = await self.vrt.compute_async(
+            self._prepare_inputs(inputs),
+            dynamic_attributes=dynamic_attributes,
+        )
+        return self._prepare_outputs(res)
+
+    def __del__(self):
+        try:
+            self.executor.shutdown()
+            self.op_executor.shutdown()
+        except Exception:
+            pass
+
+
+def _cleanup(model_path):
+    # print("Cleaning up temporary model files at", model_path)
+    rmtree(model_path)

fnnx/node_instance.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from fnnx.ops._base import BaseOp
+from typing import TypedDict
+
+
+class IO(TypedDict):
+    dtype: str
+    shape: list[int | str]
+
+
+@dataclass
+class OpInstance:
+    operator: BaseOp
+    input_specs: IO
+    output_specs: IO

fnnx/ops/_base.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+from fnnx.device import DeviceConfig
+from fnnx.dtypes import DtypesManager
+from concurrent.futures._base import Executor
+
+
+class BaseOp(ABC):
+    supported_dynamic_attributes: list[str] = []
+    required_dynamic_attributes: list[str] = []
+
+    def __init__(
+        self,
+        artifact_path: str,
+        *args,
+        attributes: dict,
+        dynamic_attribute_map: dict,
+        device_config: DeviceConfig,
+        input_specs,
+        output_specs,
+        dtypes_manager: DtypesManager,
+        executor: Executor,
+        **kwargs,
+    ):
+        self.dynamic_attribute_map = dynamic_attribute_map
+        self._warmed_up = False
+        self.artifact_path = artifact_path
+        self._device_config: DeviceConfig = device_config
+        self.attributes = attributes
+        self.input_specs = input_specs
+        self.output_specs = output_specs
+        self.dtypes_manager = dtypes_manager
+        self.executor = executor
+
+    @abstractmethod
+    def warmup(self, *args, **kwargs) -> BaseOp:
+        pass
+
+    @abstractmethod
+    def compute(self, inputs: list, dynamic_attributes: dict, **kwargs):
+        pass
+
+    @abstractmethod
+    async def compute_async(self, inputs: list, dynamic_attributes: dict, **kwargs):
+        pass
+
+    def extract_dynamic_attribute(self, dynamic_attributes: dict):
+
+        extracted = {}
+        for key, value in self.dynamic_attribute_map.items():
+            source_name = value.get("name")
+            default_value = value.get("default_value")
+            source_value = dynamic_attributes.get(source_name, None)
+            target_value = source_value or default_value
+            extracted[key] = target_value
+        return extracted
+
+    def verify_required_dynamic_attributes(self, dynamic_attributes_map: dict):
+        for key in self.required_dynamic_attributes:
+            if key not in dynamic_attributes_map:
+                raise ValueError(f"Missing required dynamic attribute: {key}")
+
+
+@dataclass
+class OpOutput:
+    value: list[Any]
+    metadata: dict | None = None

fnnx/ops/onnx.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+from fnnx.ops._base import BaseOp, OpOutput
+from os.path import join as pjoin
+
+try:
+    import onnxruntime as ort
+except ImportError:
+    ort = None
+
+try:
+    from onnxruntime_extensions import get_library_path as _get_extensions_library_path  # type: ignore
+except ImportError:
+    _get_extensions_library_path = None
+
+from fnnx.utils import to_thread
+from fnnx.device import DeviceConfig
+from fnnx.dtypes import DtypesManager
+from concurrent.futures._base import Executor
+
+CPU_EXECUTION_PROVIDER = "CPUExecutionProvider"
+CUDA_EXECUTION_PROVIDER = "CUDAExecutionProvider"
+
+
+class OnnxOp_V1(BaseOp):
+
+    def warmup(
+        self,
+    ) -> OnnxOp_V1:
+        self.model_path = pjoin(self.artifact_path, "model.onnx")
+        if not ort:
+            raise ImportError("onnxruntime is not installed")
+        if self._device_config.accelerator == "cuda":
+            execution_providers = [CUDA_EXECUTION_PROVIDER, CPU_EXECUTION_PROVIDER]
+        else:
+            execution_providers = [CPU_EXECUTION_PROVIDER]
+        session_options = ort.SessionOptions()
+        if self.attributes.get("use_onnxruntime_extensions", False):
+            if not _get_extensions_library_path:
+                raise ImportError("onnxruntime_extensions is not installed")
+            libpath = _get_extensions_library_path()
+            session_options.register_custom_ops_library(libpath)
+        self._sess = ort.InferenceSession(
+            self.model_path, providers=execution_providers, sess_options=session_options
+        )
+        self._ort_inputs = [i.name for i in self._sess.get_inputs()]
+        self._ort_outputs = [o.name for o in self._sess.get_outputs()]
+        self._warmed_up = True
+        return self
+
+    def compute(self, inputs: list, dynamic_attributes: dict, **kwargs):
+        if not self._warmed_up:
+            raise RuntimeError("Op is not warmed up")
+        outputs = self._sess.run(self._ort_outputs, dict(zip(self._ort_inputs, inputs)))
+        return OpOutput(value=list(outputs), metadata={})
+
+    async def compute_async(self, inputs: list, dynamic_attributes: dict, **kwargs):
+        return await to_thread(self.executor, self.compute, inputs, dynamic_attributes)

fnnx/registry.py

@@ -0,0 +1,23 @@
+from typing import Type
+from fnnx.ops._base import BaseOp
+from fnnx.ops.onnx import OnnxOp_V1
+import warnings
+
+
+class Registry:
+
+    def __init__(self):
+        self.ops: dict[str, Type[BaseOp]] = {}
+
+    def register_op(self, op: Type[BaseOp], name: str):
+        self.ops[name] = op
+
+    def get_op(self, name: str) -> Type[BaseOp]:
+        return self.ops[name]
+
+    def register_default_ops(self):
+        if len(self.ops.keys()) > 0:
+            warnings.warn(
+                "Attempting to register default ops into a non-empty registry."
+            )
+        self.register_op(OnnxOp_V1, "ONNX_v1")

fnnx/runtime.py

@@ -0,0 +1,35 @@
+from fnnx.handlers._base import BaseHandler
+from fnnx.handlers.local import LocalHandler, LocalHandlerConfig
+from fnnx.device import DeviceMap
+from typing import Any, Type
+
+
+class Runtime:
+    def __init__(
+        self,
+        bundle_path: str,
+        handler: Type[BaseHandler] | None = None,
+        handler_config: Any = None,
+        device_map: str | DeviceMap | None = None,
+        cleanup: bool = True,
+    ):
+        self.cleanup = cleanup
+        if handler is None:
+            handler = LocalHandler
+            handler_config = (
+                handler_config
+                if isinstance(handler_config, LocalHandlerConfig)
+                else LocalHandlerConfig()
+            )
+
+        if device_map is None:
+            device_map = DeviceMap(accelerator="cpu", node_device_map={})
+        elif isinstance(device_map, str):
+            device_map = DeviceMap(accelerator=device_map, node_device_map={})
+        self.handler: BaseHandler = handler(bundle_path, device_map, handler_config)
+
+    def compute(self, inputs: dict, dynamic_attributes: dict):
+        return self.handler.compute(inputs, dynamic_attributes)
+
+    async def compute_async(self, inputs: dict, dynamic_attributes: dict):
+        return await self.handler.compute_async(inputs, dynamic_attributes)