pipekit-experiment

Summary¶

Decision. pipekit-experiment stays the thin boundary package the scoping report and its README describe: protocols, a content-addressed registry, JSON artifacts and one adapter per tool, with nothing from those tools imported by pipekit pipelines. The registry’s design stands. Five workstreams finish the package.

• A. The registry. load cannot rebuild a model whose weights live in weights.bin (A1). A model stored without weights hashes on its class name, so two architectures collide (A2). Tags can move silently through store(name=...) and leave no history (A3). There is no gc or delete (A4); the fsspec backend needs parent directories and fails on an empty list (A5); and the scoping report’s MLflowModelRegistry does not exist (A6).

• B. Trackers. There is no ExperimentTracker implementation at all. Add a file tracker for tests and offline runs (B1), then MLflow (B2) and W&B (B3), and fix the training artifact’s tracker_run_id, which is None for the shipped Run and repr(handle) for anything else (B4).

• C. Artifacts. TrainingArtifact.rerun does not exist (C1); InferenceArtifact accepts a tag where it should pin a hash (C2); and the training artifact’s model_registry_uri omits the registry it points at (C3).

• D. Orchestration adapters. Nested pipelines do not survive HydraConfigLoader.to_yaml and from_yaml (D1, Report 9’s proposal A); MetaflowStepAdapter.as_step builds a step Metaflow rejects (D2); the DVC pointer is parsed line by line (D3).

• E. Run records. Run, RunMetrics and RunArtifacts exist, but nothing creates them, RunMetrics has no to_dataframe, and RunArtifacts.attach takes only a URI string (E1).

End state. A trained model is stored with its weights under a hash that identifies its architecture and weights, reloaded by hash or by tag in one call, promoted with a recorded tag move, and pinned by hash in an inference artifact. Its training run is logged to any tracker that satisfies the protocol, and its training artifact names the run, the data, the registry and enough of the loop to retrain.

Ownership rules¶

pipekit-experiment owns the records a run leaves behind and the registry models live in; the external tools own scheduling, dashboards and storage (Figure 1).

pipekit-experiment owns the ExperimentTracker and ModelRegistry protocols, the run records, the registries, the artifacts and the tool adapters. Its core depends only on pipekit; each tool arrives through its own extra. The registry stores weight bytes and never interprets them: each model operator (JaxModelOp, EquinoxModelOp) owns its format through serialize_weights and from_registry.

Neighbours (pipekit-train, the external tools, serving) follow the stack rules. One rule is this page’s: pipekit-train’s fallback _LocalArtifact (loop.py:174), which _build_artifact returns when pipekit-experiment is not installed (loop.py:427-433), is the only second artifact type, is never persisted, and is documented as such; removing it (making [experiment] a hard dependency of run()) is an open question.

What ships¶

Table 1 lists the public surface; the italic rows were planned and never built.

The running example: one emulator from training to an artifact¶

The stages follow the linear emulator of Report 11, a learned $W \in \mathbb{R}^{40 \times 40}$ that replaces Report 10’s transport operator $M$, through everything pipekit-experiment does to a model after training. For readable outputs some stages use a two-parameter operator Scale(factor, offset), $x \mapsto \text{factor} \cdot x + \text{offset}$, alongside it.

Stage 1 · Store the emulator under its content hash (pipekit-jax → pipekit-experiment)¶

TL;DR. Store the trained emulator and check that its address changes exactly when the model does.

Problem. The registry addresses a model by the hash (1),

where state is the operator’s class and config and $w$ its serialised weights (empty if none). Two stores of the same model must give the same $h$; changing the config or a weight must change it.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Shapes: W (40, 40), b (40,) float32 → weights.bin 6,816 bytes
import pipekit_experiment as px
from pipekit_jax import JaxModelOp

reg = px.LocalModelRegistry("registry")                                # today
op = JaxModelOp(eqx.nn.Linear(40, 40, key=key))      # W and b      today
h = reg.store(op, name="lin40", weights=op.serialize_weights())       # today
# registry/<h>/operator.json  {"class": "JaxModelOp",
#                              "config": {"module_class": "Linear"}, …}
# registry/<h>/weights.bin    6,816 bytes
# registry/<h>/metadata.json  {"tags": {}}
# registry/_tags/lin40        <h>
# same op twice → same h; one config value or one weight changed → new h
reg.store(JaxModelOp(eqx.nn.Linear(40, 40, key=k1))) == \
    reg.store(JaxModelOp(eqx.nn.Linear(3, 7, key=k2)))   # → True: no weights

The address is exactly as the code defines it, and it behaves. Writes are atomic on the local backend (a .tmp file, then os.replace), and refs are checked as a single path segment, so reg.load("../x") raises. One case fails: without weights, a JaxModelOp’s state is just {"module_class": "Linear"}, so a 40 × 40 and a 3 × 7 layer share a hash, and the second store is merged into the first one’s entry.

Stage 2 · Load it back (pipekit-experiment → pipekit-cycle)¶

TL;DR. Reload the emulator by hash and by name, with its weights, and drop it into a cycle.

Problem. A round trip must be exact: $\max \lvert W_{\text{loaded}} - W \rvert = 0$.

1
2
3
4
5
6
7
8
# Shapes: W (40, 40)
reg.load(h)
# → RuntimeError: from_state cannot reconstruct JaxModelOp: it is marked
#   forbid_in_yaml … Rebuild it in code instead.                      today
skeleton = JaxModelOp(eqx.nn.Linear(40, 40, use_bias=False, key=key0))
loaded = skeleton.from_registry(reg, "lin40")             # by tag     today
# → max |W_loaded − W| = 0.0;  a (20, 40) skeleton raises a shape error
emulator = pc.NeuralForward(loaded, dt=1.0)               # Report 10  today

ModelRegistry.load(ref) rebuilds an operator with Operator.from_state and never touches weights.bin. That works for operators whose config is plain values (Scale), and fails for every model with weights, because JaxModelOp and EquinoxModelOp are forbid_in_yaml. The path that works is the model’s own: build a skeleton of the right shape, then from_registry(reg, ref) reads the weights into it. Workstream A1 makes load do this.

Stage 3 · Promote it (pipekit-experiment)¶

TL;DR. Mark the emulator as production, then promote a candidate, and see what the registry remembers.

Problem. Deployments resolve the tag production to a hash. A promotion must be deliberate, must not be silent, and must leave a record of what production meant before.

1
2
3
4
5
6
7
8
reg.tag(h3, "production")                                            # today
reg.tag(h4, "candidate")                                             # today
reg.tag(h4, "production")
# → FileExistsError: Tag 'production' already exists. Pass force=True …
reg.tag(h4, "production", force=True)       # moves the tag             today
reg.resolve_tag("production")               # → h4; h3 is mentioned nowhere
reg.store(op5, name="production")           # moves it again, no force   today
reg.list(tags={"v": "1"})                   # → []: the stored tag is the int 1

Write-once tags are the code’s default, as the scoping report asked, and force=True is the explicit override. Two things undo that. store(..., name=...) moves a tag without force, so a script can repoint production by storing a model under that name. And no tag move is recorded: after a promotion, nothing says what production used to be. list(tags=...) compares with ==, so a string "1" does not match a stored integer 1. Workstream A3 fixes the first two.

Stage 4 · Pin it in an inference artifact (pipekit-experiment)¶

TL;DR. Freeze an inference pipeline that uses the emulator, so it can be rebuilt with exactly this model later.

Problem. An inference artifact maps each model in a pipeline to the hash it used: $\{\text{operator path} \mapsto h\}$. Rebuilding must load the same $h$, whatever the tags say by then.

1
2
3
4
5
6
art = px.InferenceArtifact(pipeline_yaml=..., backend_info={}, deps_lock=...,
                           pinned_model_hashes={"steps.1": "production"})  # today
art.save("forecast.inference.json"); px.InferenceArtifact.load(...)   # equal
art.reload_models(reg)
# → {"steps.1": Scale(factor=3.0)}: "production" resolved to whatever it is now
# with a JaxModelOp pinned → RuntimeError (from_state, as in stage 2)

pinned_model_hashes takes any ref, and reload_models resolves tags, so an artifact pinned to "production" drifts the next time the tag moves: the opposite of a pin. It also inherits stage 2’s failure for models with weights. Workstream C2 accepts only hashes.

Stage 5 · Log the training run, and read its artifact (pipekit-train → pipekit-experiment)¶

TL;DR. Train the emulator with two trackers attached, and check what the training artifact records.

Problem. One training run should reach several trackers with identical metric streams, and its artifact should name the tracker run, the data, the stored model and the registry it lives in.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Shapes: 200 pairs; 50 steps
class MemTracker:          # satisfies px.ExperimentTracker: 4 methods     user
    def start_run(self, name, config): ...
    def log_metrics(self, run, metrics, step): ...
    def log_artifact(self, run, path, name): ...
    def end_run(self, run, status): ...

a, b = MemTracker(), MemTracker()
loop = pt.TrainingLoop(..., callbacks=(pt.LogToExperiment(a),
                                       pt.LogToExperiment(b),
                                       pt.Checkpoint(registry=reg)))      # today
trained, art = loop.run()                       # art: px.TrainingArtifact
# a and b: identical mse streams, 50 points, 3.28 → 0.0029; both "completed"
art.tracker_run_id        # → None
art.model_registry_uri    # → "registry://d120fe9b…": no registry root
reg.load_weights(art.trained_model_hash)        # → None: no weights stored
pk.loads(art.training_pipeline_yaml)
# → RuntimeError: from_state cannot reconstruct TrainingLoop … forbid_in_yaml

Stacking trackers works as the scoping report promised: each LogToExperiment gets the same stream. But there is no tracker to stack; every one is user code. The artifact the loop fills has four problems, and only one is pipekit-experiment’s own:

• tracker_run_id is never the run id. _build_artifact (loop.py:414-419) keeps a callback’s run only if it has __getitem__, and then stores repr(run). The shipped Run is a dataclass without __getitem__, so the field is None, as here; a tracker that returns a mapping or a string gets its whole repr ("{'name': 'emu', 'id': 'run-0'}" in Report 11’s run), never the id alone (B4).

• model_registry_uri is registry://<hash>. It does not say which registry, and the registries give Checkpoint nothing to say it with: LocalModelRegistry exposes root, S3ModelRegistry exposes uri, and only the private _where() is common to both (C3).

• The weights are not stored. The Equinox adapter re-wraps the trained module as EquinoxModelOp, which has no serialize_weights; this is Report 11’s workstream B1.

• The pipeline YAML cannot be loaded. TrainingLoop is forbid_in_yaml, so the artifact cannot retrain (C1).

Stage 6 · Hand the pipeline to orchestration tools (pipekit-experiment → Hydra, DVC, Metaflow)¶

TL;DR. Express the pipeline as Hydra YAML, version its data with DVC and run it as a Metaflow step.

Problem. Each adapter must round-trip what it is given: a pipeline to YAML and back, a file to a DVC hash, a pipeline to a runnable step.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from pipekit_experiment.adapters.hydra import HydraConfigLoader as HL
HL.from_yaml(HL.to_yaml(Scale(2.0, 1.0)))(3.0)          # → 7.0          today
seq = pk.Sequential([Scale(2.0, 1.0), Scale(0.5, 0.0)])
HL.from_yaml(HL.to_yaml(seq))
# → InstantiationException: Sequential.operators[0] must be an Operator,
#   got DictConfig  (to_yaml wrote {class: Scale, config: …}, no _target_)

dvc = DVCDatasetVersioning(repo="dvcrepo")
dvc.add("data/train.csv"); dvc.hash("data/train.csv")   # → e5ebd4c0…    today
dvc.push("data/train.csv")      # → CalledProcessError: no remote

MetaflowStepAdapter.as_step(seq, name="scaled", inputs=["x"])
# → Metaflow: "Step scaled is missing a self.next() transition"
MetaflowStepAdapter.run(seq, self, name="scaled", inputs=["x"])  # inside a
# hand-written @step → 3.5                                              today

Flat operators round-trip through Hydra; nested ones do not, because to_yaml writes pipekit’s debug config instead of _target_ entries, which is Report 9’s proposal A seen from this side (D1). The DVC adapter works and shells out to the dvc CLI; its hash reads the first md5: or hash: line of the pointer file, so a pointer with hash: md5 first would return the string "md5" (D3). as_step builds a step without self.next, which Metaflow refuses in any position but the end, while run inside a hand-written step works (D2).

The fsspec registry was checked the same way: S3ModelRegistry("memory://") gives the same hashes and layout as the local one, but list() on an empty registry raises FileNotFoundError, a file:// registry fails on its first store unless storage_options={"auto_mkdir": True}, and its writes are not atomic (A5). Real S3 was not tested.

The scoping report’s names, checked¶

Table 2 checks the stack names the scoping report’s examples use.

Current-state audit¶

Of sixteen pieces checked, eight ship and work, six ship with defects, and two are missing (Table 3); Figure 2 counts the page’s calls by marker. The audit read references/geostack/pipekit at ef3c17f on 13 October 2026.

Smaller findings:

• Concurrent writers can race. store reads and rewrites metadata.json with no lock, and tag checks that a tag is free before writing it, so two processes can both pass the check.

• Re-storing merges tags. Storing the same model again with new tags merges them into metadata.json rather than replacing them.

• b"" and None weights hash the same, which is the intended reading of “no weights”.

• LogToExperiment never calls log_artifact, so a tracker never sees the stored model or the artifact file.

• A directory’s DVC hash carries a .dir suffix (1df8f132….dir), which a caller comparing against a file hash must expect.

Workstream A · The registry¶

Decision. The registry’s content hash and storage layout stand; what changes is that a stored model comes back whole, its address cannot collide, and every tag move is deliberate and recorded (Table 4).

1
2
3
4
5
6
7
8
9
10
11
# Shapes: W (40, 40)
# Today:
loaded = skeleton.from_registry(reg, "lin40")            # works       today
reg.load("lin40")                                        # RuntimeError
reg.store(op5, name="production")                        # silent move today

# After:
loaded = reg.load("lin40", like=skeleton)                # proposed (A1)
reg.store(op5, name="production")        # → FileExistsError  proposed (A3)
reg.tag(h5, "production", force=True)    # → appends to _history  proposed (A3)
reg.tag_history("production")            # → [(t₁, None, h3), (t₂, h3, h4), …]

Workstream B · Trackers¶

Decision. Ship trackers, starting with one that needs no service, so that LogToExperiment has something to log to out of the box (Table 5).

Workstream C · Artifacts¶

Decision. An artifact pins by hash, names its registry, and keeps enough of its training loop to run it again (Table 6).

Workstream D · Orchestration adapters¶

Decision. Each adapter round-trips what it is given; the three that ship get their gaps closed before more are added (Table 7).

Workstream E · Run records¶

Decision. The run records are produced by the trackers (B1) and given the conveniences the scoping report specified (Table 8).

Testing gates¶

Each workstream merges only when its gates pass; the existing gates are checks the code passes today (Table 9).

Sequencing and risks¶

The registry and artifact fixes come first, because every model that goes through the registry today comes back without its weights or loses its pin; trackers come second, round trips last (Figure 3).

• Write-once store(name=) breaks scripts. Scripts that store under production to promote would start raising. Fallback: warn for one release, then raise, and document tag(..., force=True) as the promotion call.

• Two writers on a shared registry. Neither tags nor metadata are locked, and S3 has no rename. Fallback: conditional writes where the backend has them (S3 If-None-Match), a lock file elsewhere, and a documented single-writer rule until then.

• Retraining needs more than a config. A TrainingLoop holds Python objects (datasets, callbacks) that YAML cannot carry. Fallback: rerun takes a builder function or a Hydra config rather than parsing training_pipeline_yaml, and checks the rebuilt dataset_hash.

• Tracker APIs churn. MLflow and W&B change often. Fallback: pin versions in the extras and run B2 and B3 against local, offline modes in CI.

• like= asks the caller for the architecture. A user loading by tag may not know the shapes. Fallback: store the module’s constructor arguments in operator.json when the model class can provide them, so load can build the skeleton itself.

Open questions¶

1. Should the registry hash include provenance? The code identifies the model by state and weights; the scoping report wanted training config and data too. This page keeps the code’s choice and links provenance through metadata (C). Is a second, provenance-inclusive hash needed for regulatory use?

2. Where does a like skeleton come from at serving time? A server that loads production needs the architecture. Is it in the deployment config, in operator.json, or in a model card stored beside the weights?

3. Does the stack need pipekit.repro.Artifact? Two JSON dataclasses cover training and inference. Is a core base class still wanted for other snapshots (use case 9’s regulatory artifact)?

4. Do tags carry stages? The scoping report used candidate and production as plain tags. Should stages be a first-class field with allowed transitions, as MLflow has?

5. Should pipekit-train keep _LocalArtifact? It lets run() work without pipekit-experiment, at the cost of a second artifact type that cannot be saved or reloaded. The alternative is to require the [experiment] extra for run() and return only TrainingArtifact.

6. A preview with each model. The scoping report’s preview.json held example inputs and outputs for sanity checks. Is it worth adding, and does it belong in the registry or the model card?

Upcoming work¶

Table 10 lists the follow-ups.

References¶

Sources¶

• The scoping report, “Report 12 — pipekit-experiment: experiment tracking and model registry” (supplied by the author).

• pipekit-experiment 0.0.2 in references/geostack/pipekit/packages/pipekit-experiment: its source, tests, README, CHANGELOG and pyproject.toml.

• The running example’s scripts and outputs, described in Section ??.