mixle.task.artifact module

Durable, portable artifacts for local task models.

mixle.utils.serialization round-trips pure probabilistic models as registry-keyed JSON, but a task model is usually torch-backed (a distilled Transformer, an MLP head), and its parameters are weights, not a JSON-serializable state. Worse, the causal LM ties head.weight = tok.weight; a naive tensor dump rejects the shared storage. This module is the missing piece: a self-describing directory that pairs

• manifest.json – how to rebuild the module (a registered builder name + its config) plus task I/O and free-form metadata, and

• weights.safetensors – the parameters, written through safetensors.torch.save_model so tied weights survive,

so a fitted model survives the process that made it. save_module/load_module are the torch path; save_json/load_json are the fallback for a pure mixle distribution. A builder is any (**config) -> nn.Module callable registered by name (register_builder); the two native architectures (mixle.causal_lm, mixle.mlp) self-register on first use, and a caller can register its own.

The acceptance bar is a fresh-process round trip: save here, load in a new interpreter from the manifest alone, get bit-identical outputs. mixle.task.model.TaskModel builds the callable task surface on top of this.

register_builder(name, builder)

Register builder under name so an artifact carrying builder=name can reconstruct its module.

builder(**config) must return a fresh (untrained) nn.Module whose parameter shapes match the saved weights. Re-registering the same name with the same callable is a no-op; a conflicting one raises.

Parameters:

• name (str)

• builder (Callable[[...], Any])

Return type:

None

get_builder(name)

Look up a registered builder, triggering native-builder self-registration on first call.

Parameters:

name (str)

Return type:

Callable[[…], Any]

class TaskManifest(payload, builder=None, config=<factory>, task='', io=<factory>, meta=<factory>, schema_version='2', created_at='', payload_sha256='', integrity_sha256='')

Bases: object

The self-describing header of a task artifact: enough to rebuild and call the model, plus provenance.

Parameters:

• payload (str)

• builder (str | None)

• config (dict[str, Any])

• task (str)

• io (dict[str, Any])

• meta (dict[str, Any])

• schema_version (str)

• created_at (str)

• payload_sha256 (str)

• integrity_sha256 (str)

to_dict()

Return the strict-JSON manifest representation written to manifest.json.

Return type:

dict[str, Any]

classmethod from_dict(d)

Parse a manifest dictionary into a TaskManifest, rejecting anything this code can’t safely read.

A manifest.json is data a fresh process trusts blindly to decide how to rebuild and load a model, so this validates it before trusting any field: artifact_type must mark it as a mixle task artifact (not some other format that happens to share a directory layout), schema_version must be one this code was actually written against – an unrecognized version (e.g. from a newer or older mixle) is rejected outright rather than silently misread as today’s schema, payload must be a payload kind this module knows how to load, and a payload that needs a builder (torch/arrays) must name one.

Parameters:

d (dict[str, Any])

Return type:

TaskManifest

read_manifest(path)

Read only the manifest of an artifact directory without loading weights.

Parameters:

path (str)

Return type:

TaskManifest

save_module(path, module, builder, config, *, task='', io=None, meta=None)

Persist a torch module as an artifact directory and return path.

builder/config must reconstruct an architecturally identical module (get_builder(builder)(**config)); weights go through safetensors.torch.save_model so tied parameters (e.g. the LM’s tied head) round-trip.

Weights and manifest are staged together and published to path in one atomic step (see _atomic_artifact_write): calling this on an existing artifact – an update – either fully replaces it or, on any failure, leaves the previous model/manifest/provenance fully intact. It never leaves new weights paired with the old manifest, or the reverse.

Parameters:

• path (str)

• module (Any)

• builder (str)

• config (dict[str, Any])

• task (str)

• io (dict[str, Any] | None)

• meta (dict[str, Any] | None)

Return type:

str

load_module(path, *, device='cpu')

Rebuild a torch module from its manifest alone and load weights; return (module, manifest).

Parameters:

• path (str)

• device (str)

Return type:

tuple[Any, TaskManifest]

save_json(path, model, *, task='', io=None, meta=None)

Persist a pure (torch-free) mixle distribution via the safe serialization registry; return path.

model.json and the manifest are staged together and published to path in one atomic step (see _atomic_artifact_write), so an update that fails partway – e.g. a non-serializable replacement model – never leaves path with a new model paired with the old manifest, or a truncated file of either kind.

Parameters:

• path (str)

• model (Any)

• task (str)

• io (dict[str, Any] | None)

• meta (dict[str, Any] | None)

Return type:

str

load_json(path)

Rebuild a pure mixle distribution from a json-payload artifact; return (model, manifest).

Parameters:

path (str)

Return type:

tuple[Any, TaskManifest]

register_arrays_builder(name, builder)

Register builder(arrays: dict[str, ndarray], **config) -> model for arrays-payload artifacts.

The arrays payload is for torch-free numeric students (e.g. an int8-quantized MLP): weights live in one .npz, and the builder reconstructs the runnable model from them in a fresh process.

Parameters:

• name (str)

• builder (Callable[[...], Any])

Return type:

None

get_arrays_builder(name)

Look up a registered arrays builder, triggering native self-registration on first call.

Parameters:

name (str | None)

Return type:

Callable[[…], Any]

save_arrays(path, arrays, builder, config=None, *, task='', io=None, meta=None)

Persist a dict of numpy arrays as an artifact directory (arrays.npz); return path.

Arrays and manifest are staged together and published to path in one atomic step (see _atomic_artifact_write), so a failed update never leaves new arrays paired with the old manifest.

Written compressed. This payload holds quantized edge students whose whole point is bytes on a device, and every member costs a fixed ~128-byte .npy header plus zip framing – on a small student the per-member framing outweighs the weights themselves. np.load reads compressed and stored archives identically, so this is transparent to load_arrays() and to any artifact written by an earlier version.

Parameters:

• path (str)

• arrays (dict[str, Any])

• builder (str)

• config (dict[str, Any] | None)

• task (str)

• io (dict[str, Any] | None)

• meta (dict[str, Any] | None)

Return type:

str

load_arrays(path)

Rebuild a torch-free model from an arrays-payload artifact; return (model, manifest).

Parameters:

path (str)

Return type:

tuple[Any, TaskManifest]