Skip to content

Latest commit

 

History

History
890 lines (593 loc) · 25.6 KB

api.md

File metadata and controls

890 lines (593 loc) · 25.6 KB

Balsa API reference

Regenerate with pixi run -e docs docs; verify with pixi run -e docs docs-check. The raw compiler reference is written to build/balsa-api.json. Compiler signatures may expand default values and normalize type names.

All entries below are available from balsa. This is the supported public surface; other declarations in implementation modules are internal. See the 0.2 migration guide for storage, validation, ownership and compatibility conventions.

Index

decode

Defined in balsa.packed.

def decode[dtype: DType = DType.float32](var data: List[UInt8], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> PackedModel[dtype]

Own checkpoint bytes with bounds checks; semantic validation is opt-in.

def decode[origin: ImmOrigin, //, dtype: DType = DType.float32](data: Span[UInt8, origin], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> PackedModel[dtype]

Copy borrowed input once so the packed result owns its bytes.

load

Defined in balsa.packed.

def load[dtype: DType = DType.float32](path: String, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> PackedModel[dtype]

Load directly into owned packed storage; no editable forest is allocated.

decode_auto

Defined in balsa.packed.

def decode_auto(var data: List[UInt8], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> AnyPackedModel

Discover precision and own the original checkpoint buffer.

def decode_auto[origin: ImmOrigin, //](data: Span[UInt8, origin], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> AnyPackedModel

Discover precision and copy borrowed bytes into an independent owner.

load_auto

Defined in balsa.packed.

def load_auto(path: String, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> AnyPackedModel

PackedModel

Defined in balsa.packed.

struct PackedModel[dtype: DType]

Own a checkpoint; field views cannot outlive or mutate its storage.

Traits: AnyType, Deinitable, Movable.

__init__

def __init__(out self, var data: List[UInt8], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises

num_trees

def num_trees(self) -> Int

num_features

def num_features(self) -> Int

postprocessor_name

def postprocessor_name(self) raises -> String

attributes_text

def attributes_text(self) raises -> String

base_scores

def base_scores(self) -> Span[Float64, origin_of(self._header.base_scores)]

target_ids

def target_ids(self) -> Span[Int32, origin_of(self._header.target_id)]

class_ids

def class_ids(self) -> Span[Int32, origin_of(self._header.class_id)]

tree

def tree(self, index: Int) raises -> PackedTree[dtype, origin_of(self._data)]

Borrow typed fields; tree lookup scans its field headers, not payloads.

to_model

def to_model(self, options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=True)) raises -> Model[dtype]

Materialize independent editable fields, validating by default.

validate

def validate(self, options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises

Check semantics serially using one reusable tree and topology scratch.

Always checks, even with enabled=False. max_workers is a cap; this representation currently uses one worker. No editable forest is built.

AnyPackedModel

Defined in balsa.packed.

comptime AnyPackedModel = Variant[PackedModel[DType.float32], PackedModel[DType.float64]]

encode

Defined in balsa.storage.

def encode[dtype: DType](model: PackedModel[dtype]) -> List[UInt8]

Copy retained serialization without semantic revalidation.

def encode(model: Variant[PackedModel[DType.float32], PackedModel[DType.float64]]) -> List[UInt8]

Copy retained serialization without semantic revalidation.

def encode[dtype: DType](model: Model[dtype], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> List[UInt8]

Serialize editable fields; pass enabled=True options for semantic checks.

def encode(model: Variant[Model[DType.float32], Model[DType.float64]], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> List[UInt8]

Serialize editable fields; pass enabled=True options for semantic checks.

save

Defined in balsa.storage.

def save[dtype: DType](model: PackedModel[dtype], path: String) raises

Copy retained serialization without semantic revalidation.

def save(model: Variant[PackedModel[DType.float32], PackedModel[DType.float64]], path: String) raises

Copy retained serialization without semantic revalidation.

def save[dtype: DType](model: Model[dtype], path: String, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises

Serialize editable fields; pass enabled=True options for semantic checks.

def save(model: Variant[Model[DType.float32], Model[DType.float64]], path: String, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises

Serialize editable fields; pass enabled=True options for semantic checks.

decode_editable

Defined in balsa.codec.

def decode[dtype: DType = DType.float32](var data: List[UInt8], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> Model[dtype]

Decode editable fields; automatic semantic validation is opt-in.

Wire bounds, format checks and parser limits always apply. Precision must match the file tags. A skipped validation pass can be run with validate().

def decode[origin: ImmOrigin, //, dtype: DType = DType.float32](data: Span[UInt8, origin], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> Model[dtype]

Borrow checkpoint bytes and return independently owned model fields.

decode_into

Defined in balsa.codec.

def decode_into[dtype: DType, origin: ImmOrigin](mut model: Model[dtype], data: Span[UInt8, origin], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises

Decode into owned storage, reusing existing field capacities.

On failure fields may be partially updated and must not be used as a valid model. The destination remains safe to destroy or pass to decode_into again. Input bytes are borrowed only for this call. All parser checks still apply.

load_editable

Defined in balsa.codec.

def load[dtype: DType = DType.float32](path: String, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> Model[dtype]

Load a checkpoint of the specified precision (float32 by default).

decode_auto_editable

Defined in balsa.codec.

def decode_auto(var data: List[UInt8], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> AnyModel

Discover precision from the checkpoint, preserving its typed storage.

def decode_auto[origin: ImmOrigin](data: Span[UInt8, origin], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> AnyModel

Discover precision from borrowed bytes and return owned model fields.

load_auto_editable

Defined in balsa.codec.

def load_auto(path: String, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises -> AnyModel

Load either supported precision without converting its values.

checkpoint_dtype

Defined in balsa.codec.

def checkpoint_dtype(data: List[UInt8], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024))) raises -> DType

Inspect the v4 header and both dtype tags; does not validate the body.

def checkpoint_dtype[origin: ImmOrigin](data: Span[UInt8, origin], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024))) raises -> DType

Inspect precision directly from borrowed checkpoint bytes.

AnyModel

Defined in balsa.codec.

comptime AnyModel = Variant[Model[DType.float32], Model[DType.float64]]

Owned model with discovered float32 or float64 precision; accepted by encode and save.

Operator

Defined in balsa.constants.

struct Operator

Numerical comparison codes; NONE is reserved for non-split nodes.

Traits: AnyType, Deinitable, Movable.

NONE

comptime NONE = Int8(0)

EQ

comptime EQ = Int8(1)

LT

comptime LT = Int8(2)

LE

comptime LE = Int8(3)

GT

comptime GT = Int8(4)

GE

comptime GE = Int8(5)

NodeType

Defined in balsa.constants.

struct NodeType

Node kind codes for scalar/vector leaves and numerical/categorical splits.

Traits: AnyType, Deinitable, Movable.

LEAF

comptime LEAF = Int8(0)

NUMERICAL

comptime NUMERICAL = Int8(1)

CATEGORICAL

comptime CATEGORICAL = Int8(2)

TaskType

Defined in balsa.constants.

struct TaskType

Treelite task codes preserved as metadata; Balsa does not run inference.

Traits: AnyType, Deinitable, Movable.

BINARY_CLF

comptime BINARY_CLF = UInt8(0)

REGRESSOR

comptime REGRESSOR = UInt8(1)

MULTI_CLF

comptime MULTI_CLF = UInt8(2)

LEARNING_TO_RANK

comptime LEARNING_TO_RANK = UInt8(3)

ISOLATION_FOREST

comptime ISOLATION_FOREST = UInt8(4)

TypeInfo

Defined in balsa.constants.

struct TypeInfo

Wire tags. UINT32 exists in the enum but is unsupported in v4 models.

Traits: AnyType, Deinitable, Movable.

INVALID

comptime INVALID = UInt8(0)

UINT32

comptime UINT32 = UInt8(1)

FLOAT32

comptime FLOAT32 = UInt8(2)

FLOAT64

comptime FLOAT64 = UInt8(3)

Model

Defined in balsa.model.

struct Model[dtype: DType]

Owned model fields, validated by the codec.

Traits: AnyType, Copyable, Deinitable, Movable.

Field Type Description
major Int32 Treelite checkpoint major version, independent of Balsa version.
minor Int32 Treelite checkpoint minor version.
patch Int32 Treelite checkpoint patch version.
threshold_type UInt8 TypeInfo tag matching the model dtype.
leaf_output_type UInt8 TypeInfo tag matching the model dtype.
num_tree UInt64 Number of trees; must equal the trees list length.
num_feature Int32 Number of input features.
task_type UInt8 TaskType code preserved as model metadata.
average_tree_output UInt8 Wire boolean indicating averaged tree output.
num_target Int32 Number of output targets.
num_class List[Int32] Class count for each target.
leaf_vector_shape List[Int32] Two dimensions describing a vector leaf output.
target_id List[Int32] Per-tree target annotation; -1 denotes all targets.
class_id List[Int32] Per-tree class annotation; -1 denotes all classes.
postprocessor List[UInt8] Postprocessor name bytes; Balsa never executes the postprocessor.
sigmoid_alpha Float32 Stored postprocessor coefficient.
ratio_c Float32 Stored postprocessor coefficient.
base_scores List[Float64] Base score array with num_target * max(num_class) entries.
attributes List[UInt8] Opaque attribute JSON bytes; not parsed or validated as JSON.
extensions List[Extension] Opaque records in the model extension slot.
trees List[Tree[dtype]] Owned trees in serialization order.

Scalar

comptime Scalar = Scalar[dtype]

__init__

def __init__(out self)

Create raw storage with checkpoint version 4.6.1 and matching dtype tags.

The empty metadata is incomplete; use ModelBuilder for valid defaults.

postprocessor_name

def postprocessor_name(self) raises -> String

Decode UTF-8 strictly; unknown names are allowed and bytes untouched.

set_postprocessor_name

def set_postprocessor_name(mut self, name: String)

Store UTF-8 text without restricting postprocessor names.

attributes_text

def attributes_text(self) raises -> String

Decode UTF-8 strictly; does not parse or validate attribute JSON.

set_attributes_text

def set_attributes_text(mut self, text: String)

Store text; the caller remains responsible for valid attribute JSON.

Tree

Defined in balsa.model.

struct Tree[dtype: DType]

Owned tree fields, validated by the codec.

Traits: AnyType, Copyable, Deinitable, Movable.

Field Type Description
num_nodes Int32 Number of nodes; node zero is the root.
has_categorical_split UInt8 Wire boolean indicating whether the tree contains categorical splits.
node_type List[Int8] Per-node NodeType code.
cleft List[Int32] Per-node left child index; -1 for leaves.
cright List[Int32] Per-node right child index; -1 for leaves.
split_index List[Int32] Per-node feature index for split nodes.
default_left List[UInt8] Per-node wire boolean selecting the left branch for missing values.
leaf_value List[Scalar[dtype]] Per-node scalar leaf value.
threshold List[Scalar[dtype]] Per-node numerical split threshold.
cmp List[Int8] Per-node Operator code for numerical splits.
category_list_right_child List[UInt8] Per-node wire boolean: category membership selects the right child.
leaf_vector List[Scalar[dtype]] Concatenated vector-leaf values.
leaf_vector_begin List[UInt64] Per-node inclusive offset into leaf_vector.
leaf_vector_end List[UInt64] Per-node exclusive offset into leaf_vector.
category_list List[UInt32] Concatenated categorical split IDs.
category_list_begin List[UInt64] Per-node inclusive offset into category_list.
category_list_end List[UInt64] Per-node exclusive offset into category_list.
data_count List[UInt64] Optional per-node sample counts, paired with data_count_present.
data_count_present List[UInt8] Presence flags for data_count; both arrays may be empty.
sum_hess List[Float64] Optional per-node Hessian sums, paired with sum_hess_present.
sum_hess_present List[UInt8] Presence flags for sum_hess; both arrays may be empty.
gain List[Float64] Optional per-node gains, paired with gain_present.
gain_present List[UInt8] Presence flags for gain; both arrays may be empty.
tree_extensions List[Extension] Opaque records in the tree extension slot.
node_extensions List[Extension] Opaque records in the node extension slot.

Scalar

comptime Scalar = Scalar[dtype]

__init__

def __init__(out self)

Create empty tree storage; use TreeBuilder for a complete tree.

leaf_values

def leaf_values(self, node: Int) raises -> Span[Scalar[dtype], origin_of(self.leaf_vector)]

Borrow a node's vector leaf payload (empty for a scalar leaf).

Checks node and offset bounds even after raw field edits. The view is immutable and cannot outlive the tree or coexist with its mutation.

categories

def categories(self, node: Int) raises -> Span[UInt32, origin_of(self.category_list)]

Borrow a node's category payload, checking node and offset bounds.

Extension

Defined in balsa.model.

struct Extension

Opaque optional field retained in its original extension slot.

Traits: AnyType, Copyable, Deinitable, Movable.

Field Type Description
name List[UInt8] Opaque field name bytes.
element_size UInt64 Bytes per payload element.
count UInt64 Number of payload elements; payload length must equal count * element_size.
payload List[UInt8] Opaque bytes preserved without interpretation.

ModelBuilder

Defined in balsa.builder.

struct ModelBuilder[dtype: DType = DType.float32]

Build a model with derived tree/target counts and default scalar metadata.

Traits: AnyType, Deinitable, Movable.

__init__

def __init__(out self, num_feature: Int32, *, task_type: UInt8 = TaskType.REGRESSOR, average_tree_output: Bool = False, var num_class: List[Int32] = List(Int32(1), __list_literal__=NoneType(None)), var leaf_vector_shape: List[Int32] = List(Int32(1), Int32(1), __list_literal__=NoneType(None)), var base_scores: List[Float64] = List(), postprocessor: String = "identity", sigmoid_alpha: Float32 = Float32("1"), ratio_c: Float32 = Float32("1"), attributes: String = "", limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024))) raises

Initialize and validate metadata for a model without trees.

num_class defines one entry per target. Empty base_scores become zeros with num_target * max(num_class) entries. List arguments are consumed. Postprocessor settings are stored as metadata and never executed.

add_tree

def add_tree(mut self, var tree: Tree[dtype], *, target_id: Int32 = Int32(0), class_id: Int32 = Int32(0)) raises

Transfer tree ownership; annotation and topology checks run at build.

build

def build(deinit self, options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=True)) raises -> Model[dtype]

Consume the builder, validating the model unless options disables it.

TreeBuilder

Defined in balsa.builder.

struct TreeBuilder[dtype: DType = DType.float32]

Fill dense node IDs in any order; node 0 is the root.

Each node must be defined exactly once. build() checks completeness; ModelBuilder.build() checks topology, features and output shapes.

Traits: AnyType, Deinitable, Movable.

Scalar

comptime Scalar = Scalar[dtype]

__init__

def __init__(out self, num_nodes: Int, limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024))) raises

Allocate a tree with positive, bounded node count and undefined nodes.

leaf

def leaf(mut self, node: Int, value: Scalar[dtype]) raises

Define a scalar leaf without parallel-array bookkeeping.

leaf_vector

def leaf_vector(mut self, node: Int, values: List[Scalar[dtype]]) raises

Define a vector leaf; shape is checked against model metadata.

numerical_split

def numerical_split(mut self, node: Int, *, feature: Int32, threshold: Scalar[dtype], left: Int32, right: Int32, op: Int8 = Operator.LT, default_left: Bool = False) raises

Define a numerical split; default_left chooses the missing-value branch.

categorical_split

def categorical_split(mut self, node: Int, *, feature: Int32, categories: List[UInt32], left: Int32, right: Int32, default_left: Bool = False, categories_right: Bool = False) raises

Define a categorical split, copying category IDs into tree storage.

Membership selects the right child when categories_right is True, otherwise the left child. default_left selects the missing-value branch.

build

def build(deinit self) raises -> Tree[dtype]

Consume the builder, rejecting undefined nodes; no array copies.

validate

Defined in balsa.validation.

def validate[dtype: DType](model: Model[dtype], limits: Limits = Limits(Int(67108864), Int(8388608), Int(100000), Int(1000000), Int(1024)), options: ValidationOptions = ValidationOptions(Int(4), Int(4096), Int(65536), enabled=False)) raises

Reject malformed metadata, lengths, topology and segment offsets.

Always validates, including when options.enabled is False.

ValidationOptions

Defined in balsa.validation.

struct ValidationOptions

Control automatic model validation and its concurrency per call.

Codec calls skip semantic validation by default (enabled=False). Pass enabled=True for untrusted or uncertain inputs. Builders and packed to_model() explicitly default to checked construction/conversion. Explicit validate() always checks the model, regardless of enabled. Thresholds apply together. The shared runtime may use fewer workers.

Traits: AnyType, Copyable, Deinitable, Movable.

Field Type Description
enabled Bool Whether codec and builder calls validate automatically; default False.
max_workers Int Per-call worker cap; values <= 1 select serial validation.
min_trees Int Minimum tree count for parallel validation; default 4,096.
min_nodes Int Minimum total node count for parallel validation; default 65,536.

__init__

def __init__(out self, max_workers: Int = Int(4), min_trees: Int = Int(4096), min_nodes: Int = Int(65536), *, enabled: Bool = False)

Set automatic validation policy and parallel execution thresholds.

max_workers <= 1 selects serial execution. Both thresholds must be met and no tree may contain more than half the nodes to use parallelism.

Limits

Defined in balsa.wire.

struct Limits

Allocation/work limits for untrusted checkpoints (not an RSS guarantee).

Traits: AnyType, Copyable, Deinitable, Movable.

Field Type Description
max_bytes Int Maximum checkpoint size in bytes; default 64 MiB.
max_elements Int Maximum elements in an individual array; default 8,388,608.
max_trees Int Maximum trees per model; default 100,000.
max_nodes Int Maximum total nodes per model; default 1,000,000.
max_extensions Int Maximum records in each extension slot; default 1,024.

__init__

def __init__(out self, max_bytes: Int = Int(67108864), max_elements: Int = Int(8388608), max_trees: Int = Int(100000), max_nodes: Int = Int(1000000), max_extensions: Int = Int(1024))

Set byte, per-array, tree, total-node and per-slot extension caps.

validate

def validate(self) raises

Reject limits outside the inclusive range 1 through 2^30.

__version__

Defined in balsa.

comptime __version__ = "0.2.0"

Balsa library version, independent of the Treelite checkpoint version.

version

Defined in balsa.

def version() -> String

Return the library version; prefer the compile-time version constant.

Retained for compatibility with existing callers.