"""Example 10 — A nested pipeline: model 1 → model 2 → two producer options, compared. Some training pipelines are staged: an upstream model labels the data, a downstream model trains on those labels and emits an intermediate signal, and only then does a final decoding step turn that signal into an answer — often with more than one decoding **option** to choose between. This wires that whole shape together in datavo and scores the two final options. It also shows the workset/dataset choice in practice: the **base** model trains on a stable **dataset** (like 09, the reproducible foundation), while the **follow-up** stages train and produce on cheap, live **worksets** — training on a workset is fine when you don't need to reproduce the exact source, and a workset is a first-class **Add** source too (``from_workset``). base dataset ──▶ model 1 ──produce teacher.cls──▶ base store (model 1 generates a key) │ workset pins teacher.cls = model 1 ▼ model 2 (trains on a workset) ──produce score.cls──▶ base store (intermediate) │ Add(score.cls, label.cls) into ONE separate store ▼ final store ┌───────────────┴───────────────┐ optA ─produce decision.cls optB ─produce decision.cls └───────────────┬───────────────┘ accuracy(decision|optA) ◀ compare ▶ accuracy(decision|optB) The datavo mechanics that make each arrow work were each introduced earlier; here they compose: * **Chained produce.** model 1 attaches ``teacher.cls`` under producer ``model1``; model 2's follow-up training **workset** pins it (``producer_pins={"teacher.cls": "model1"}``, see 06) so the downstream stage consumes exactly that upstream output. model 2 then attaches ``score.cls``. * **Move an intermediate into its own store.** An ``Add`` from a **workset** (worksets are first-class Add sources, ``from_workset``) brings ``score.cls`` into a fresh ``final`` store as an ordinary, readable key (02) — the substrate the two options fan out over. * **Two options, one key.** ``optA`` and ``optB`` are two producers writing the **same** ``decision.cls`` on the one final store (two decode strategies over the one intermediate). Two producers sharing a key in a single store is the 06 shape, so the two outcomes are compared the 06 way — one **workset** per option, each **pinning** its producer — and scored against the label. (A plain read can't separate them: two producers write the same key, so telling them apart needs the pinned read.) All model/decoder logic is a deterministic stand-in (no real training); the pinning, chaining, Add, and pinned reads are the real datavo operations. DATAVO_SERVER=https://dev.datavo.io python examples/10_nested_pipeline_compare.py """ from __future__ import annotations import tempfile from pathlib import Path from _common import build_dataset, cleanup, cleanup_datasets, make_client, seed_source_store, step, unique from datavo_sdk import ( SourceSpec, SampleStream, HasKeys, KeyGroupSpec, AttachedBy, Not, ShardPlan, TarWriter, ) # Ground-truth labelled audio (N divisible by 6 keeps the stand-in fractions exact). N = 24 SAMPLES = [ {"__key__": f"{i:04d}", "audio.wav": f"audio-{i}".encode(), "label.cls": i % 2} for i in range(N) ] def read_scalars(client, store: str, sample_ids: list[str], keys: list[str]) -> dict[str, dict[str, int]]: """Read the given scalar ``keys`` per sample id straight from the store (server-side ``get_sample`` attribute reads — no dataset/streaming). Each scalar key is a sample attribute with a ``numeric_value``; a stage reads its inputs this way to compute outputs.""" out: dict[str, dict[str, int]] = {} for sid in sample_ids: attrs = {a.name: a.numeric_value for a in client.get_sample(sid).attributes} out[sid] = {key: int(attrs[key]) for key in keys if attrs.get(key) is not None} return out def missing_ids(client, store: str, producer_id: str, key: str, *, inputs: list[str]) -> list[str]: """The ids in ``store`` that still need ``key`` from ``producer_id`` (03_attach_keys.py's shape): an attach workset whose row filter subtracts what this producer already attached, iterated for its ids. The id set *is* the work, so doneness is an empty workset.""" workset = client.create_workset( sources=[SourceSpec( sample_store=store, keys=[inputs[0]], filters=[ HasKeys(inputs), Not(AttachedBy(producer_id, [key])), ], )], shard_plan=ShardPlan(shard_size=8), ) if not workset.sample_count: return [] # decode=False: a producer forwards payloads rather than interpreting them. return [ sample["__key__"] for sample in SampleStream.from_workset( workset.workset_id, keys=[inputs[0]], client=client, decode=False ) ] def attach_scalar(client, store: str, producer_id: str, key: str, value_by_id: dict[str, int]) -> None: """Attach one scalar ``key`` under ``producer_id`` for the ids in ``value_by_id`` (condensed 03_attach_keys.py). An empty map is an idempotent no-op.""" if not value_by_id: return with tempfile.TemporaryDirectory() as tmp: shard = Path(tmp) / f"{producer_id}-{key}.tar" with TarWriter(str(shard)) as writer: for sid, value in value_by_id.items(): writer.write({"__key__": sid, key: value}) client.attach_keys( store, str(shard), keys=[key], producer_id=producer_id, ) def graded_score(label: int, rank: int, group_size: int) -> int: """Stand-in for model 2's intermediate confidence in {0,1,2} from model 1's pseudo-label. Positives are mostly *weak* (score 1) with a strong (score 2) minority; negatives are mostly *clear* (score 0) with a borderline (score 1) minority. This graded middle is what makes the two decode thresholds below actually disagree — and lets ground truth pick a winner. """ minority = rank < group_size // 3 # first third of each class is the "confident" minority if label == 1: return 2 if minority else 1 return 1 if minority else 0 def main() -> None: client = make_client() base = unique("example_pipe_base") ds1 = unique("example_pipe_ds1") # the base model's dataset — the one stable, reusable source final = unique("example_pipe_final") eval_worksets: dict[str, str] = {} # per-option worksets, filled in stage 4 try: step(f"seed the first (labelled) store {base!r} with {N} samples; the BASE model trains on a dataset {ds1!r}") # The base model gets a stable, reproducible source (a dataset, like 09). The follow-up # stages below iterate on cheap worksets — training on a workset is fine when you don't # need to keep/reproduce the exact source. seed_source_store(client, base, samples=SAMPLES) build_dataset( client, ds1, sources=[SourceSpec(sample_store=base, keys=["audio.wav", "label.cls"])], key_groups=[KeyGroupSpec(name="audio", keys=["audio.wav"]), KeyGroupSpec(name="label", keys=["label.cls"])], shard_plan=ShardPlan(shard_size=8), ) # ── Stage 1 — model 1 trains on ds1 and GENERATES a key (teacher.cls) ────────────── step("model 1 (trains on ds1) → produce teacher.cls (a pseudo-label per sample)") ids = missing_ids(client, base, "model1", "teacher.cls", inputs=["audio.wav"]) base_labels = read_scalars(client, base, ids, ["label.cls"]) # A clean teacher for the demo: its pseudo-label equals the ground-truth label. teacher = {sid: vals["label.cls"] for sid, vals in base_labels.items()} client.register_producer(base, "model1", declared_output_keys=["teacher.cls"], metadata={"producer_kind": "teacher"}) attach_scalar(client, base, "model1", "teacher.cls", teacher) # ── Stage 2 — model 2 CONSUMES model 1's key, produces an intermediate ───────────── step("model 2's follow-up training runs on a WORKSET (fast, exploratory), pinning teacher.cls = model1") # No need to freeze a dataset to iterate: a workset is directly trainable — you just # don't get a reproducible artifact, which is fine for a follow-up run. ws2 = client.create_workset( sources=[SourceSpec( sample_store=base, keys=["audio.wav", "teacher.cls", "label.cls"], producer_pins={"teacher.cls": "model1"}, # consume model 1's output specifically )], key_groups=[KeyGroupSpec(name="audio", keys=["audio.wav"]), KeyGroupSpec(name="teacher", keys=["teacher.cls"]), KeyGroupSpec(name="label", keys=["label.cls"])], shard_plan=ShardPlan(shard_size=8), ).workset_id client.wait_for_workset(ws2, timeout_seconds=600) teacher_rows = list(SampleStream.from_workset(ws2, keys=["teacher.cls"], client=client)) teacher_mean = sum(int(r["teacher.cls"]) for r in teacher_rows) / len(teacher_rows) print(f" the follow-up workset resolves model 1's pinned key: mean(teacher.cls) = {teacher_mean:.2f}") step("model 2 (trains on ds2) → produce score.cls (intermediate confidence in {0,1,2})") client.register_producer(base, "model2", declared_output_keys=["score.cls"], metadata={"producer_kind": "student"}) # model 2 has its own missing set, and reads its input (model 1's key) back from the store. score_ids = missing_ids(client, base, "model2", "score.cls", inputs=["audio.wav", "teacher.cls"]) pseudo = {sid: v["teacher.cls"] for sid, v in read_scalars(client, base, score_ids, ["teacher.cls"]).items()} # Rank within each class so the confident minority is deterministic (no per-id randomness). pos = sorted(sid for sid, v in pseudo.items() if v == 1) neg = sorted(sid for sid, v in pseudo.items() if v == 0) rank = {sid: r for r, sid in enumerate(pos)} | {sid: r for r, sid in enumerate(neg)} score = {sid: graded_score(pseudo[sid], rank[sid], len(pos) if pseudo[sid] == 1 else len(neg)) for sid in pseudo} attach_scalar(client, base, "model2", "score.cls", score) # ── Stage 3 — move the intermediate into ONE separate store ──────────────────────── step(f"select the intermediate as a WORKSET (score.cls pinned to model2 + label), Add it into {final!r}") # A workset is a first-class Add source too — no need to freeze a dataset just to move a # selection into another store. ws_int = client.create_workset( sources=[SourceSpec( sample_store=base, keys=["score.cls", "label.cls"], producer_pins={"score.cls": "model2"}, )], key_groups=[KeyGroupSpec(name="score", keys=["score.cls"]), KeyGroupSpec(name="label", keys=["label.cls"])], shard_plan=ShardPlan(shard_size=8), ).workset_id client.wait_for_workset(ws_int, timeout_seconds=600) client.create_store(name=final, display_name="final decode store") client.add_to_store(final, operation_id="seed-final", from_workset=ws_int, key_subset=["score.cls", "label.cls"], display_name="intermediate results") client.wait_for_add(final, "seed-final", timeout_seconds=300) # In `final` the added score.cls is an ordinary readable key — the two options fan out over it. # ── Stage 4 — two producer OPTIONS turn the intermediate into a final decision ───── step("two decode options write the SAME decision.cls on the one final store") # Two thresholds on the graded score: optA is lenient (score ≥ 1 ⇒ 1), optB is strict # (score ≥ 2 ⇒ 1). They disagree exactly on the graded-middle samples. Register both up # front (declared_output_keys are fixed at registration) before the first attach. options = {"optA": 1, "optB": 2} for option, threshold in options.items(): client.register_producer(final, option, declared_output_keys=["decision.cls"], metadata={"producer_kind": "decoder", "threshold": threshold}) for option, threshold in options.items(): # Each option has its own missing set: idempotency is per (producer_id, key), so the # second option sees every sample as its own work even though the key name is shared. final_ids = missing_ids(client, final, option, "decision.cls", inputs=["score.cls"]) final_scores = read_scalars(client, final, final_ids, ["score.cls"]) decision = {sid: int(final_scores[sid]["score.cls"] >= threshold) for sid in final_ids} attach_scalar(client, final, option, "decision.cls", decision) # One workset per option, pinning ITS decision.cls next to the ground-truth label. eval_worksets[option] = client.create_workset( sources=[SourceSpec( sample_store=final, keys=["decision.cls", "label.cls"], producer_pins={"decision.cls": option}, )], key_groups=[KeyGroupSpec(name="decision", keys=["decision.cls"]), KeyGroupSpec(name="label", keys=["label.cls"])], shard_plan=ShardPlan(shard_size=8), ).workset_id # ── Compare the two final outcomes vs the metric ─────────────────────────────────── step("compare: open each option (pin-aware) and score decision.cls against the label") results: dict[str, tuple[float, int]] = {} for option in options: rows = list(SampleStream.from_workset(eval_worksets[option], keys=["decision.cls", "label.cls"], client=client)) hits = sum(int(int(r["decision.cls"]) == int(r["label.cls"])) for r in rows) results[option] = (hits / len(rows), len(rows)) print(f" {'option':<7} {'threshold':>9} {'accuracy':>9} {'n':>3}") for option, threshold in options.items(): acc, n = results[option] print(f" {option:<7} {threshold:>9} {acc:>8.1%} {n:>3}") winner = max(results, key=lambda option: results[option][0]) print(f" → {winner} decodes the intermediate better " f"({results[winner][0]:.1%} vs {min(a for a, _ in results.values()):.1%})") finally: # Teardown: the final store (the Add's target) → the base model's dataset (pins base) # → base store. The follow-up worksets (ws2, ws_int, and the per-option eval worksets) # are ephemeral — an unused one is released on its own, and a workset never blocks deletion. cleanup(client, final) cleanup_datasets(client, ds1) cleanup(client, base) if __name__ == "__main__": main()