Merge pull request #9 from BeastByteAI/pyfnx_utils

added model reader/writer, bumped version to 0.0.8

Author: OKUA1

spec/schemas/src/generate.py

@@ -1,5 +1,7 @@
 import json
 import subprocess
+import os
+import shutil
 from pydantic_models.op_instances import OpInstances
 from pydantic_models.manifest import Manifest
 from pydantic_models.variants.pipeline import PipelineVariant
@@ -73,3 +75,38 @@
     f.write("# This file is auto generated and must not be modified manually!\n")
     f.write(f"schema = {combined_json_schema}")
 subprocess.run(["black", spec_py_path])
+
+
+HEADER_TEXT = (
+    "# ==============================================================\n"
+    "# This file was automatically copied from spec.\n"
+    "# DO NOT EDIT — changes here will be overwritten.\n"
+    "# ==============================================================\n\n"
+)
+
+
+def copy_with_header(src_dir, dst_dir, header_text=HEADER_TEXT):
+    for root, _, files in os.walk(src_dir):
+        rel_path = os.path.relpath(root, src_dir)
+        dst_root = os.path.join(dst_dir, rel_path)
+        os.makedirs(dst_root, exist_ok=True)
+
+        for f in files:
+            if f.endswith(".py"):
+                src_path = os.path.join(root, f)
+                dst_path = os.path.join(dst_root, f)
+
+                with open(src_path, "r", encoding="utf-8") as s:
+                    original = s.read()
+
+                with open(dst_path, "w", encoding="utf-8") as d:
+                    d.write(HEADER_TEXT + original)
+            elif f.endswith(".pyc"):
+                pass
+            else:
+                shutil.copy2(os.path.join(root, f), os.path.join(dst_root, f))
+
+
+copy_with_header(
+    "./pydantic_models", "./../../../src/python/fnnx/extras/pydantic_models"
+)

src/python/fnnx/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.7"
+__version__ = "0.0.8"

src/python/fnnx/extras/builder.py

@@ -0,0 +1,301 @@
+import tarfile
+import io
+from typing import Type, Callable
+from fnnx.variants.pyfunc import PyFunc
+import json
+from dataclasses import dataclass
+from dataclasses import asdict as dataclass_asdict
+import inspect
+import sys
+
+from pydantic import BaseModel as PydanticBaseModel
+
+from fnnx import __version__ as fnnx_version
+
+from fnnx.extras.pydantic_models.manifest import Manifest, NDJSON, JSON, Var
+from fnnx.extras.pydantic_models.variants.pyfunc import PyFuncVariant
+from fnnx.extras.pydantic_models.envs import Python3_CondaPip, PipDependency
+
+
+def asdict(obj):
+    if isinstance(obj, PydanticBaseModel):
+        return obj.model_dump()
+    return dataclass_asdict(obj)
+
+
+def asjson(obj):
+    if isinstance(obj, PydanticBaseModel):
+        return obj.model_dump_json(indent=4)
+    return json.dumps(asdict(obj), indent=4)
+
+
+PYTHON_VERSION = (
+    f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+)
+
+
+@dataclass
+class PyFuncSpec:
+    filepath: str
+    class_name: str
+
+
+class PyfuncBuilder:
+    def __init__(
+        self,
+        pyfunc: Type[PyFunc] | PyFuncSpec,
+        model_name: str | None = None,
+        model_version: str | None = None,
+        model_description: str | None = None,
+        create_meta_callback: Callable | None = None,
+    ) -> None:
+        self._inputs: list[NDJSON | JSON] = []
+        self._outputs: list[NDJSON | JSON] = []
+        self._dynamic_attributes: list[Var] = []
+        self._env_vars: list[Var] = []
+
+        self._producer_name: str = "fnnx.ai"
+        self._producer_version: str = fnnx_version
+        self._producer_tags: list[str] = []
+
+        self._extra_dtypes: dict = {}
+
+        self._name = model_name
+        self._version = model_version
+        self._description = model_description
+
+        self._extra_modules = []
+        self._extra_files = []
+        self._extra_values: dict | None = None
+
+        self._build_dependencies = []
+        self._rt_dependencies = []
+
+        self.create_meta_callback = create_meta_callback
+
+        if isinstance(pyfunc, PyFuncSpec):
+            self.pyfunc_name = pyfunc.class_name
+            pyfunc_file = pyfunc.filepath
+        elif issubclass(pyfunc, PyFunc):
+            self.pyfunc_name = pyfunc.__name__
+            pyfunc_file = inspect.getfile(pyfunc)
+        else:
+            raise TypeError(
+                "Pyfunc must be a subclass of PyFunc or an instance of PyFuncSpec"
+            )
+
+        with open(pyfunc_file) as f:
+            self.pyfunc_content = f.read()
+
+    def add_input(self, input_spec: NDJSON | JSON) -> None:
+        if not isinstance(input_spec, (NDJSON, JSON)):
+            raise TypeError("input_spec must be NDJSON or JSON instance")
+        if input_spec.name in [x.name for x in self._inputs]:
+            raise ValueError(f"input with name {input_spec.name} already exists")
+        if (
+            input_spec.dtype.startswith("ext::")
+            and input_spec.dtype not in self._extra_dtypes
+        ):
+            raise ValueError(f"extra dtype with name {input_spec.dtype} not defined")
+        self._inputs.append(input_spec)
+
+    def add_output(self, output_spec: NDJSON | JSON) -> None:
+        if not isinstance(output_spec, (NDJSON, JSON)):
+            raise TypeError("output_spec must be NDJSON or JSON instance")
+        self._outputs.append(output_spec)
+
+    def add_dynamic_attribute(self, dynamic_attribute: Var) -> None:
+        if not isinstance(dynamic_attribute, Var):
+            raise TypeError("dynamic_attribute must be Var instance")
+        if dynamic_attribute.name in [x.name for x in self._dynamic_attributes]:
+            raise ValueError(
+                f"dynamic_attribute with name {dynamic_attribute.name} already exists"
+            )
+        self._dynamic_attributes.append(dynamic_attribute)
+
+    def add_env_var(self, env_var: Var) -> None:
+        if not isinstance(env_var, Var):
+            raise TypeError("env_var must be Var instance")
+        if env_var.name in [x.name for x in self._env_vars]:
+            raise ValueError(f"env_var with name {env_var.name} already exists")
+        self._env_vars.append(env_var)
+
+    def set_extra_values(self, values: dict) -> None:
+        self._extra_values = values.copy()
+
+    def define_dtype(self, name: str, dtype: Type[PydanticBaseModel]) -> None:
+        if not name.startswith("ext::"):
+            raise ValueError("dtype name must start with 'ext::'")
+        self._extra_dtypes[name] = dtype.model_json_schema()
+
+    def set_producer_info(
+        self, name: str, version: str, tags: list[str] | None = None
+    ) -> None:
+        self._producer_name = name
+        self._producer_version = version
+        self._producer_tags = tags or []
+
+    def add_module(self, module_path: str) -> None:
+        module_name = module_path.split("/")[-1]
+        if module_name in [x.split("/")[-1] for x in self._extra_modules]:
+            raise ValueError(f"module with name {module_name} already exists")
+        self._extra_modules.append(module_path)
+
+    def add_file(self, file_path: str, target_path: str) -> None:
+        if target_path in [x[1] for x in self._extra_files]:
+            raise ValueError(f"file with name {target_path} already exists")
+        self._extra_files.append((file_path, target_path))
+
+    def save(self, path: str) -> None:
+        manifest = Manifest(
+            variant="pyfunc",
+            producer_name=self._producer_name,
+            producer_version=self._producer_version,
+            producer_tags=self._producer_tags,
+            inputs=self._inputs,
+            outputs=self._outputs,
+            dynamic_attributes=self._dynamic_attributes,
+            env_vars=self._env_vars,
+            name=self._name,
+            version=self._version,
+            description=self._description,
+        )
+
+        f = File(path)
+        f.create_file("manifest.json", asjson(manifest))
+        f.create_file("dtypes.json", json.dumps(self._extra_dtypes))
+        f.create_file("env.json", json.dumps(self._make_env()))
+        f.create_file("variant_artifacts/__pyfunc__.py", self.pyfunc_content)
+        f.create_file(
+            "variant_config.json",
+            asjson(
+                PyFuncVariant(
+                    pyfunc_classname=self.pyfunc_name, extra_values=self._extra_values
+                )
+            ),
+        )
+        f.create_file("ops.json", "[]")
+        f.make_artifacts_folders()
+        if self.create_meta_callback:
+            self.create_meta_callback(f)
+        else:
+            f.create_file("meta.json", "[]")
+
+        for module in self._extra_modules:
+            f.copy(
+                module,
+                f"variant_artifacts/extra_modules/{module.split('/')[-1]}",
+                should_exclude_pycache=True,
+            )
+
+        for file_path, target_path in self._extra_files:
+            f.copy(file_path, f"variant_artifacts/extra_files/{target_path}")
+
+        try:
+            pass
+        except Exception as e:
+            raise e
+        finally:
+            f.close()
+
+    def add_default_build_dependencies(self) -> None:
+        import subprocess
+        import re
+
+        def get_version(
+            command: str, idx: int = 0, regex: str = r"^\d+(\.\d+)*([a-zA-Z]+\d*)?$"
+        ) -> str:
+            try:
+                result = subprocess.run(
+                    command,
+                    shell=True,
+                    check=True,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                )
+                output = result.stdout.strip()
+
+                version_parts = output.split()
+                version = version_parts[idx]
+                if not version or not re.match(regex, version):
+                    raise ValueError(f"Invalid version format: ```{version}```")
+                return version
+            except (subprocess.CalledProcessError, IndexError, ValueError) as e:
+                print(f"Error retrieving version for '{command}': {e}")
+                return "unknown"
+
+        try:
+            pip_version = get_version("pip --version", 1)
+            self.add_build_dependency(f"pip=={pip_version}")
+        except Exception as e:
+            print(f"Error adding pip version to build dependencies: {e}")
+
+        try:
+            setuptools_version = get_version(
+                "pip show setuptools | grep Version | awk '{print $2}'"
+            )
+            self.add_build_dependency(f"setuptools=={setuptools_version}")
+        except Exception as e:
+            print(f"Error adding setuptools version to build dependencies: {e}")
+
+        try:
+            wheel_version = get_version(
+                "pip show wheel | grep Version | awk '{print $2}'"
+            )
+            self.add_build_dependency(f"wheel=={wheel_version}")
+        except Exception as e:
+            print(f"Error adding wheel version to build dependencies: {e}")
+
+    def add_build_dependency(self, dep: str) -> None:
+        self._build_dependencies.append(dep)
+
+    def add_runtime_dependency(self, dep: str) -> None:
+        # TODO add conditions for runtime dependencies
+        self._rt_dependencies.append(dep)
+
+    def add_fnnx_runtime_dependency(self, core: bool = False) -> None:
+        import fnnx
+
+        fnnx_version = fnnx.__version__
+        dependency_name = "fnnx[core]" if core else "fnnx"
+        self.add_runtime_dependency(f"{dependency_name}=={fnnx_version}")
+
+    def _make_env(self):
+        return {
+            "python3::conda_pip": asdict(
+                Python3_CondaPip(
+                    python_version=PYTHON_VERSION,
+                    build_dependencies=self._build_dependencies,
+                    dependencies=[PipDependency(package=p) for p in self._rt_dependencies],
+                )
+            )
+        }
+
+
+class File:
+    def __init__(self, path):
+        self.tar = tarfile.open(path, "w")
+
+    def make_artifacts_folders(self):
+        self.create_file("meta_artifacts/.keep", ".keep")
+        self.create_file("ops_artifacts/.keep", ".keep")
+        self.create_file("variant_artifacts/extra_modules/.keep", ".keep")
+        self.create_file("variant_artifacts/extra_files/.keep", ".keep")
+
+    def create_file(self, path, content: str):
+        data = io.BytesIO(content.encode())
+        info = tarfile.TarInfo(path)
+        info.size = len(data.getvalue())
+        self.tar.addfile(info, data)
+
+    def copy(self, src, dst, should_exclude_pycache: bool = False):
+        def exclude_pycache(tarinfo):
+            if "__pycache__" in tarinfo.name and should_exclude_pycache:
+                return None
+            return tarinfo
+
+        self.tar.add(src, dst, filter=exclude_pycache)
+
+    def close(self):
+        self.tar.close()