mlbundle 0.1.0a1__py3-none-any.whl

mlbundle/_composite/io.py

@@ -0,0 +1,158 @@
+"""Composite orchestration; each byte member belongs to its existing loader."""
+
+from dataclasses import dataclass
+from typing import Any
+
+from pydantic import JsonValue
+from sklearn.pipeline import Pipeline
+
+from mlbundle._native_io import NativeExport, prepare_native
+from mlbundle._pipeline_contract import PipelineInfo
+from mlbundle._sklearn import transformers
+from mlbundle._sklearn.pipeline_scope import validate_components
+from mlbundle.claims import ModelClaims, verify_claims, verify_structure
+from mlbundle.contract import Library, MemberName, SerializationFormat
+from mlbundle.declarations import ExportDeclarations
+from mlbundle.errors import LimitExceededError, UnsupportedModelError
+from mlbundle.limits import ArchiveLimits
+from mlbundle.registry import get_adapter
+
+from .claims import extract, outer_claims, rebuild, structure, verify
+from .prediction import Predictor
+
+
+@dataclass(frozen=True)
+class CompositeExport:
+    leaf: NativeExport
+    transformer_adapter: transformers.Adapter
+    prefix: Pipeline | None
+    restored_prefix: Pipeline | None
+    step_names: list[str]
+    verbose: bool
+    claims: ModelClaims
+    members: dict[MemberName, bytes]
+    library: Library = Library.SKLEARN
+    serialization: SerializationFormat = SerializationFormat.PIPELINE
+
+    @property
+    def structure(self) -> dict[str, JsonValue] | None:
+        return self.leaf.structure
+
+    def finish(self, names: list[str]) -> tuple[PipelineInfo, dict[str, JsonValue]]:
+        expected = extract(
+            self.prefix,
+            self.leaf.claims,
+            self.leaf.library,
+            self.step_names,
+            self.verbose,
+            names=names,
+            transformer_adapter=self.transformer_adapter,
+        )
+        actual_leaf = self.leaf.adapter.extract_claims(self.leaf.restored)
+        actual_model = rebuild(
+            self.restored_prefix, self.leaf.restored, self.step_names, self.verbose
+        )
+        validate_components(actual_model)
+        actual_claims = outer_claims(self.restored_prefix, actual_leaf, self.verbose)
+        verify_claims(self.claims, actual_claims)
+        actual = extract(
+            self.restored_prefix,
+            actual_leaf,
+            self.leaf.library,
+            self.step_names,
+            self.verbose,
+            names=names,
+            transformer_adapter=self.transformer_adapter,
+        )
+        verify(expected, actual)
+        original = structure(self.leaf.structure, expected)
+        encoded = verify_structure(
+            original,
+            structure(self.leaf.adapter.extract_structure(self.leaf.restored), actual),
+        )
+        self.members[MemberName.STRUCTURE] = encoded
+        return expected, original
+
+
+def prepare(
+    model: Any, declarations: ExportDeclarations, limits: ArchiveLimits
+) -> CompositeExport:
+    validate_components(model, exporting=True)
+    leaf_model = model.steps[-1][1]
+    roots = {base.__module__.split(".")[0] for base in type(leaf_model).__mro__}
+    library = Library.XGBOOST if "xgboost" in roots else Library.CATBOOST
+    adapter = get_adapter(library)
+    if not adapter.recognizes(leaf_model):
+        raise UnsupportedModelError(
+            reason="Composite final estimator must be a supported XGBoost or CatBoost wrapper."
+        )
+    leaf = prepare_native(leaf_model, adapter, declarations, limits)
+    prefix = restored = None
+    transformer_adapter = transformers.Adapter()
+    members = dict(leaf.members)
+    if len(model.steps) > 1:
+        prefix = Pipeline(
+            list(model.steps[:-1]),
+            memory=None,
+            verbose=model.verbose,
+            transform_input=None,
+        )
+        data = transformer_adapter.serialize(prefix)
+        if len(data) > limits.member_bytes(MemberName.SKOPS):
+            raise LimitExceededError("Model member size limit exceeded.")
+        members[MemberName.SKOPS] = data
+        restored = transformer_adapter.deserialize(data)
+    claims = outer_claims(prefix, leaf.claims, model.verbose)
+    return CompositeExport(
+        leaf,
+        transformer_adapter,
+        prefix,
+        restored,
+        [name for name, _ in model.steps],
+        model.verbose,
+        claims,
+        members,
+    )
+
+
+def load_components(
+    recipe: PipelineInfo,
+    members: dict[MemberName, bytes],
+    names: list[str],
+    verbose: bool,
+) -> tuple[Pipeline, Predictor, ModelClaims, dict[str, JsonValue]]:
+    adapter = get_adapter(recipe.estimator.library)
+    transformer_adapter = transformers.Adapter()
+    leaf = adapter.deserialize(members[adapter.model_member])
+    prefix = (
+        None
+        if recipe.transformers is None
+        else transformer_adapter.deserialize(members[MemberName.SKOPS])
+    )
+    native_claims = adapter.extract_claims(leaf)
+    model = rebuild(prefix, leaf, recipe.step_names, verbose)
+    validate_components(model)
+    actual = extract(
+        prefix,
+        native_claims,
+        adapter.library,
+        recipe.step_names,
+        verbose,
+        names=names,
+        transformer_adapter=transformer_adapter,
+    )
+    verify(recipe, actual)
+    claims = outer_claims(prefix, native_claims, verbose)
+    extracted = structure(adapter.extract_structure(leaf), actual)
+    settings = native_claims.hyperparameters.get("prediction")
+    categorical = (
+        adapter.library == Library.XGBOOST
+        and isinstance(settings, dict)
+        and settings.get("enable_categorical") is True
+    )
+    return (
+        model,
+        Predictor(prefix, adapter.predictor(leaf), claims, categorical=categorical),
+        claims,
+        extracted,
+    )

mlbundle/_composite/prediction.py

@@ -0,0 +1,77 @@
+"""Transform original inputs, then delegate to the cached native predictor."""
+
+from typing import Any
+
+from sklearn import config_context
+from sklearn.pipeline import Pipeline
+
+from mlbundle._prediction_input import tabular
+from mlbundle._sklearn.prediction import _string_inputs
+from mlbundle.adapters import (
+    Predictor as NativePredictor,
+    ProbabilityOutput,
+    RegressionOutput,
+)
+from mlbundle.claims import ModelClaims
+from mlbundle.contract import ModelKind
+from mlbundle.errors import InvalidModelError, MLBundleError, PredictionInputError
+
+
+class Predictor:
+    def __init__(
+        self,
+        prefix: Pipeline | None,
+        native: NativePredictor,
+        claims: ModelClaims,
+        *,
+        categorical: bool = False,
+    ) -> None:
+        self.prefix, self.native, self.claims = prefix, native, claims
+        self.strings = prefix is not None and _string_inputs(prefix)
+        self.categorical = categorical
+
+    def _transform(self, data: Any) -> Any:
+        # A single-step composite preserves the leaf's categorical admission.
+        if self.prefix is None:
+            return data
+        data, rows = tabular(
+            data,
+            self.claims.n_features,
+            strings=self.strings,
+            categorical=self.categorical,
+        )
+        names = getattr(data, "columns", None)
+        if (
+            names is not None
+            and self.claims.feature_names is not None
+            and list(names) != self.claims.feature_names
+        ):
+            raise PredictionInputError(
+                "Prediction columns must match model feature names and order."
+            )
+        try:
+            with config_context(
+                transform_output="default", enable_metadata_routing=False
+            ):
+                result = self.prefix.transform(data)
+            if len(result.shape) != 2 or result.shape[0] != rows:
+                raise InvalidModelError(
+                    "Composite transformation must preserve prediction rows."
+                )
+            return result
+        except MLBundleError:
+            raise
+        except Exception:
+            raise InvalidModelError(
+                "Composite Pipeline transformation failed."
+            ) from None
+
+    def predict(self, data: object) -> RegressionOutput:
+        if self.claims.kind != ModelKind.REGRESSOR:
+            raise ValueError("predict requires a regressor.")
+        return self.native.predict(self._transform(data))
+
+    def predict_proba(self, data: object) -> ProbabilityOutput:
+        if self.claims.kind != ModelKind.CLASSIFIER:
+            raise ValueError("predict_proba requires a classifier.")
+        return self.native.predict_proba(self._transform(data))

mlbundle/_container_rules.py

@@ -0,0 +1,73 @@
+"""Shared container names, member sets and resource checks."""
+
+from collections.abc import Iterable
+from os import PathLike
+from typing import BinaryIO
+
+from .contract import MODEL_MEMBERS, MemberName
+from .errors import InvalidContainerError, LimitExceededError
+from .limits import ArchiveLimits
+from .manifest import Manifest
+
+BundleSource = str | PathLike[str] | bytes | BinaryIO
+CHUNK_BYTES = 64 * 1024
+MEMBER_ORDER = (
+    MemberName.MANIFEST,
+    MemberName.SKOPS,
+    MemberName.UBJSON,
+    MemberName.CBM,
+    MemberName.ONNX,
+    MemberName.SCHEMA,
+    MemberName.STRUCTURE,
+    MemberName.SAMPLES,
+)
+
+
+def effective_limits(limits: ArchiveLimits | None) -> ArchiveLimits:
+    # Revalidate even model_copy/model_construct values; ceilings are mandatory.
+    return ArchiveLimits.model_validate((limits or ArchiveLimits()).model_dump())
+
+
+def member_name(raw: bytes | str) -> MemberName:
+    try:
+        return MemberName(raw if isinstance(raw, str) else raw.decode("ascii"))
+    except (UnicodeError, ValueError):
+        raise InvalidContainerError("Unknown or unsafe member name.") from None
+
+
+def check_members(
+    names: Iterable[MemberName], manifest: Manifest | None = None
+) -> None:
+    present = set(names)
+    if MemberName.MANIFEST not in present:
+        raise InvalidContainerError("Missing manifest.json.")
+    models = {pair[1] for pair in MODEL_MEMBERS.values()}
+    selected = present & models
+    if len(selected) != 1 and selected not in (
+        {MemberName.SKOPS, MemberName.UBJSON},
+        {MemberName.SKOPS, MemberName.CBM},
+    ):
+        raise InvalidContainerError("Invalid model member combination.")
+    if manifest is not None and present != {MemberName.MANIFEST, *manifest.integrity}:
+        raise InvalidContainerError("Members disagree with manifest references.")
+
+
+def check_sizes(
+    sizes: Iterable[tuple[MemberName, int, int]], limits: ArchiveLimits
+) -> None:
+    compressed = uncompressed = count = 0
+    for name, packed, unpacked in sizes:
+        count += 1
+        compressed += packed
+        uncompressed += unpacked
+        if packed > limits.max_compressed_bytes or unpacked > limits.member_bytes(name):
+            raise LimitExceededError(f"Member size limit exceeded: {name.value}.")
+        if unpacked > limits.max_compression_ratio * packed:
+            raise LimitExceededError(f"Compression ratio limit exceeded: {name.value}.")
+    if count > limits.max_members:
+        raise LimitExceededError("Member count limit exceeded.")
+    if (
+        compressed > limits.max_compressed_bytes
+        or uncompressed > limits.max_total_bytes
+    ):
+        raise LimitExceededError("Total member size limit exceeded.")

mlbundle/_export_manifest.py

@@ -0,0 +1,78 @@
+"""Manifest construction shared by standalone and composite export."""
+
+from collections.abc import Mapping, Sequence
+from datetime import UTC, datetime
+from importlib.metadata import version
+from uuid import uuid4
+
+from pydantic import JsonValue
+
+from .claims import ModelClaims
+from .contract import (
+    FORMAT_NAME,
+    FORMAT_VERSION,
+    FeatureOrigin,
+    Library,
+    MemberName,
+    SerializationFormat,
+)
+from .environment import capture_environment
+from .training import TrainingDeclaration
+
+
+def manifest_parts(
+    claims: ModelClaims,
+    library: Library,
+    serialization: SerializationFormat,
+    *,
+    structure: dict[str, JsonValue] | None,
+    feature_names: Sequence[str] | None,
+    training: TrainingDeclaration | None,
+    target_name: str | None,
+    metadata: Mapping[str, JsonValue] | None,
+) -> dict[str, object]:
+    names = claims.feature_names
+    origin = FeatureOrigin.MODEL
+    if names is None:
+        origin = (
+            FeatureOrigin.DECLARED
+            if feature_names is not None
+            else FeatureOrigin.GENERATED
+        )
+        names = (
+            list(feature_names)
+            if feature_names is not None
+            else [f"feature_{index}" for index in range(1, claims.n_features + 1)]
+        )
+        if training is not None:
+            names, origin = training.names, FeatureOrigin.DATAFRAME
+    model = claims.model_dump(
+        mode="json", exclude={"feature_names", "feature_dtypes", "class_labels", "onnx"}
+    )
+    model.update(library=library.value, serialization=serialization.value)
+    return {
+        "format_name": FORMAT_NAME,
+        "format_version": FORMAT_VERSION,
+        "bundle_id": str(uuid4()),
+        "created_at": datetime.now(UTC).isoformat(),
+        "writer_version": version("mlbundle"),
+        "model": model,
+        "features": {
+            "names": names,
+            "dtypes": [None] * claims.n_features
+            if training is None
+            else training.dtypes,
+            "model_dtypes": claims.feature_dtypes,
+            "origin": origin.value,
+        },
+        "target": {"class_labels": claims.class_labels, "name": target_name},
+        "environment": capture_environment(onnx=claims.onnx).model_dump(mode="json"),
+        "references": {
+            "structure": None if structure is None else MemberName.STRUCTURE.value,
+            "schema": None if training is None else MemberName.SCHEMA.value,
+            "samples": None
+            if training is None or training.samples is None
+            else MemberName.SAMPLES.value,
+        },
+        "user_metadata": {} if metadata is None else dict(metadata),
+    }

mlbundle/_native_io.py

@@ -0,0 +1,53 @@
+"""Shared native preparation and full round-trip verification."""
+
+from dataclasses import dataclass
+
+from pydantic import JsonValue
+
+from .adapters import Adapter
+from .claims import ModelClaims, verify_claims, verify_structure
+from .contract import MODEL_MEMBERS, Library, MemberName, SerializationFormat
+from .declarations import ExportDeclarations, resolve_labels
+from .errors import LimitExceededError
+from .limits import ArchiveLimits
+
+
+@dataclass(frozen=True)
+class NativeExport:
+    adapter: Adapter
+    model: object
+    restored: object
+    claims: ModelClaims
+    structure: dict[str, JsonValue] | None
+    members: dict[MemberName, bytes]
+
+    @property
+    def library(self) -> Library:
+        return self.adapter.library
+
+    @property
+    def serialization(self) -> SerializationFormat:
+        return MODEL_MEMBERS[self.library][0]
+
+
+def prepare_native(
+    model: object,
+    adapter: Adapter,
+    declarations: ExportDeclarations,
+    limits: ArchiveLimits,
+) -> NativeExport:
+    model = adapter.prepare_export(model, declarations)
+    data = adapter.serialize(model)
+    if len(data) > limits.member_bytes(adapter.model_member):
+        raise LimitExceededError("Model member size limit exceeded.")
+    claims = resolve_labels(adapter.extract_claims(model), declarations)
+    structure = adapter.extract_structure(model)
+    restored = adapter.deserialize(data)
+    verify_claims(
+        claims, resolve_labels(adapter.extract_claims(restored), declarations)
+    )
+    structure_bytes = verify_structure(structure, adapter.extract_structure(restored))
+    members = {adapter.model_member: data}
+    if structure is not None:
+        members[MemberName.STRUCTURE] = structure_bytes
+    return NativeExport(adapter, model, restored, claims, structure, members)

mlbundle/_onnx/__init__.py

@@ -0,0 +1,115 @@
+"""Original ONNX bytes with a validated CPU session; optional imports stay here."""
+
+from dataclasses import dataclass
+from os import PathLike
+from pathlib import Path
+from typing import Any
+
+import onnx
+from pydantic import JsonValue, ValidationError
+
+from mlbundle.claims import ModelClaims
+from mlbundle.contract import Library, MemberName
+from mlbundle.declarations import ExportDeclarations, resolve_labels
+from mlbundle.errors import InvalidModelError, LimitExceededError, MLBundleError
+from mlbundle.limits import ArchiveLimits
+
+from .description import Description, describe
+from .prediction import Predictor
+from .session import create_session
+from .validation import parse_model
+
+
+@dataclass(frozen=True)
+class OnnxModel:
+    """Validated native inference session and the exact bytes that created it."""
+
+    data: bytes
+    session: Any
+
+
+def _description(data: bytes) -> Description:
+    try:
+        return describe(parse_model(data))
+    except MLBundleError:
+        raise
+    except ValidationError as error:
+        failure = error.errors(include_input=False, include_url=False)[0]
+        # Disclose closed field names, never author-supplied keys or values.
+        fields = {"producer_name", "producer_version", "feature_names", "class_labels"}
+        field = next((part for part in failure["loc"] if part in fields), None)
+        if field is not None and failure["type"] == "string_too_long":
+            limit = failure.get("ctx", {})["max_length"]
+            raise InvalidModelError(
+                f"ONNX {field} exceeds the {limit}-character limit."
+            ) from None
+        raise InvalidModelError("Invalid ONNX model claims.") from None
+    except Exception:
+        raise InvalidModelError("Invalid ONNX model claims.") from None
+
+
+class Adapter:
+    library = Library.ONNX
+    model_member = MemberName.ONNX
+
+    def __init__(self) -> None:
+        self._data: bytes | None = None
+        self._validated: Description | None = None
+
+    def _describe(self, data: bytes) -> Description:
+        # An adapter belongs to one export/load, never a global cache or bundle.
+        if data is not self._data or self._validated is None:
+            description = _description(data)
+            self._data, self._validated = data, description
+        return self._validated
+
+    def recognizes(self, model: object) -> bool:
+        return isinstance(model, (bytes, onnx.ModelProto, OnnxModel)) or (
+            isinstance(model, (str, PathLike)) and Path(model).suffix.lower() == ".onnx"
+        )
+
+    def prepare_export(self, model: object, declarations: ExportDeclarations) -> object:
+        try:
+            if isinstance(model, OnnxModel):
+                data = model.data
+            elif isinstance(model, bytes):
+                data = model
+            elif isinstance(model, onnx.ModelProto):
+                if model.ByteSize() > ArchiveLimits().max_member_bytes:
+                    raise LimitExceededError("ONNX model exceeds the byte limit.")
+                data = model.SerializeToString()
+            else:
+                path = Path(model)  # type: ignore[arg-type]
+                cap = ArchiveLimits().max_member_bytes
+                if path.stat().st_size > cap:
+                    raise LimitExceededError("ONNX model exceeds the byte limit.")
+                with path.open("rb") as stream:
+                    data = stream.read(cap + 1)
+            description = self._describe(data)
+            resolve_labels(description.claims, declarations)
+            return self.deserialize(data)
+        except MLBundleError:
+            raise
+        except Exception:
+            raise InvalidModelError("Cannot read ONNX model input.") from None
+
+    def serialize(self, model: Any) -> bytes:
+        return model.data
+
+    def deserialize(self, data: bytes) -> OnnxModel:
+        description = self._describe(data)
+        try:
+            return OnnxModel(data, create_session(data, description))
+        except MLBundleError:
+            raise
+        except Exception:
+            raise InvalidModelError("ONNX CPU session creation failed.") from None
+
+    def extract_claims(self, model: Any) -> ModelClaims:
+        return self._describe(model.data).claims
+
+    def extract_structure(self, model: object) -> dict[str, JsonValue] | None:  # noqa: ARG002
+        return None
+
+    def predictor(self, model: Any) -> Predictor:
+        return Predictor(model.session, self._describe(model.data))

mlbundle/_onnx/_wire_runs.py

@@ -0,0 +1,103 @@
+"""Bounded C regex scans of repeated wire scalars; no protobuf objects allocated."""
+
+import re
+from functools import lru_cache
+
+import numpy as np
+
+# Match exactly the same uint64 encodings accepted by the scalar walker,
+# including overlong encodings; byte ten may contain only bit zero.
+VARINT = rb"(?:[\x00-\x7f]|[\x80-\xff]{1,8}[\x00-\x7f]|[\x80-\xff]{9}[\x00-\x01])"
+CHUNK_BYTES = 64 * 1024
+PACKED = re.compile(rb"(?:" + VARINT + rb")+")
+
+
+@lru_cache(maxsize=128)
+def _pattern(tag: bytes, wire: int, value: bytes):
+    scalar = VARINT if wire == 0 else rb".{%d}" % (8 if wire == 1 else 4)
+    unit = re.escape(tag) + scalar if wire != 2 else re.escape(value)
+    return re.compile(rb"(?:" + unit + rb")+", re.DOTALL)
+
+
+def _varint_chunk(data: bytes, start: int, end: int, tag: bytes):
+    values = np.frombuffer(data, dtype=np.uint8, count=end - start, offset=start)
+    terminals = np.flatnonzero(values < 128)
+    count = len(terminals) // 2
+    if not count:
+        return start, 0
+    pairs = terminals[: count * 2].reshape(-1, 2)
+    starts = np.concatenate(([0], pairs[:-1, 1] + 1))
+    valid = pairs[:, 0] == starts + len(tag) - 1
+    for offset, byte in enumerate(tag):
+        # Clipping is only for malformed, too-short tags; their offset check fails.
+        valid &= values[np.minimum(starts + offset, len(values) - 1)] == byte
+    widths = pairs[:, 1] - pairs[:, 0]
+    valid &= (widths <= 10) & ((widths != 10) | (values[pairs[:, 1]] <= 1))
+    invalid = np.flatnonzero(~valid)
+    count = int(invalid[0]) if len(invalid) else count
+    return start + int(pairs[count - 1, 1]) + 1 if count else start, count
+
+
+def _fixed_chunk(data: bytes, start: int, end: int, tag: bytes, wire: int):
+    size = len(tag) + (8 if wire == 1 else 4)
+    count = (end - start) // size
+    values = np.frombuffer(
+        data, dtype=np.uint8, count=count * size, offset=start
+    ).reshape(-1, size)
+    valid = np.all(values[:, : len(tag)] == np.frombuffer(tag, dtype=np.uint8), axis=1)
+    invalid = np.flatnonzero(~valid)
+    count = int(invalid[0]) if len(invalid) else count
+    return start + count * size, count
+
+
+def repeated(
+    data: bytes, start: int, end: int, *, value_start: int, stop: int, wire: int
+):
+    """Include the already validated first field, then scan equivalent wire units."""
+    if wire == 2 and stop - start > 4096:
+        return stop, 1
+    tag = data[start:value_start]
+    pattern = _pattern(tag, wire, data[start:stop] if wire == 2 else b"")
+    probe = pattern.match(data, start, min(end, start + 64))
+    vectorize = wire != 2 and probe is not None and probe.end() - start >= 32
+    chunk_bytes = 128 if vectorize else CHUNK_BYTES
+    count, position = 0, start
+    while position < end and data.startswith(tag, position, end):
+        chunk_end = min(end, position + chunk_bytes)
+        if vectorize:
+            next_position, units = (
+                _varint_chunk(data, position, chunk_end, tag)
+                if wire == 0
+                else _fixed_chunk(data, position, chunk_end, tag, wire)
+            )
+            # Grow only after proportional progress; short alternating fields
+            # never speculate over a full 64 KiB window.
+            if next_position - position >= (chunk_end - position) // 2 and units:
+                count += units
+                position = next_position
+                chunk_bytes = min(CHUNK_BYTES, 2 * chunk_bytes)
+                continue
+        match = pattern.match(data, position, chunk_end)
+        if match is None:
+            break
+        next_position = match.end()
+        if wire == 0:
+            # Every field has one terminating tag byte and one scalar byte.
+            values = np.frombuffer(
+                data, dtype=np.uint8, count=next_position - position, offset=position
+            )
+            count += int(np.count_nonzero(values < 128)) // 2
+        else:
+            size = stop - start if wire == 2 else len(tag) + (8 if wire == 1 else 4)
+            count += (next_position - position) // size
+        position = next_position
+    return position, count
+
+
+def packed_varints(data: bytes, position: int, end: int):
+    match = PACKED.match(data, position, min(end, position + CHUNK_BYTES))
+    if match is None:
+        return position, 0
+    stop = match.end()
+    values = np.frombuffer(data, dtype=np.uint8, count=stop - position, offset=position)
+    return stop, int(np.count_nonzero(values < 128))