Tracking an ML experiment with ancestree¶

This notebook walks through a small but complete machine-learning workflow and uses ancestree to record the lineage of every artifact it produces.

We take the classic Iris dataset and fan it out across a grid of modelling choices:

• two scalers — RobustScaler and StandardScaler
• three embeddings — UMAP, t-SNE and PCA
• two clusterers — K-Means and HDBSCAN

That is 2 × 3 × 2 = 12 clustering results, each with its own provenance. Instead of juggling file names by hand, we let ancestree capture how each node was derived from its parent, so the whole experiment stays queryable once it has finished running.

1. Imports¶

The usual scientific-Python stack, plus ancestree itself.

In [1]:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.datasets import load_iris
from sklearn.preprocessing import RobustScaler, StandardScaler
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans, HDBSCAN
from sklearn.metrics import calinski_harabasz_score
import umap

import ancestree

import numpy as np import pandas as pd import matplotlib.pyplot as plt from sklearn.datasets import load_iris from sklearn.preprocessing import RobustScaler, StandardScaler from sklearn.decomposition import PCA from sklearn.manifold import TSNE from sklearn.cluster import KMeans, HDBSCAN from sklearn.metrics import calinski_harabasz_score import umap import ancestree

/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

2. Configure the lineage store¶

A LineageStore is the on-disk record of the experiment. Two settings shape it:

• rules declare which step types may follow which. Each key is a step type and its value lists the permitted parents, where [None] marks a valid pipeline root. This keeps provenance honest — a clustering node can only ever hang off an embedding.
• gen_triggers name the step types that begin a new generation, which is how ancestree groups successive runs of the pipeline.

The configuration is written to .lineage_config.json inside the store, so it is picked up automatically next time — reopening with ancestree.LineageStore("ML_pipeline") is enough.

In [2]:

rules = {
    "raw data": [None],  # a pipeline starts from raw data
    "scaling": ["raw data"],  # scaling follows the raw data
    "embedding": ["scaling"],  # embeddings are computed on scaled data
    "clustering": ["embedding"],  # clusters are found in the embedding space
}
triggers = ["scaling", "embedding", "clustering"]

store = ancestree.LineageStore(
    "../assets/demo", rules=rules, gen_triggers=triggers, dedupe=True, chunk=True
)

rules = { "raw data": [None], # a pipeline starts from raw data "scaling": ["raw data"], # scaling follows the raw data "embedding": ["scaling"], # embeddings are computed on scaled data "clustering": ["embedding"], # clusters are found in the embedding space } triggers = ["scaling", "embedding", "clustering"] store = ancestree.LineageStore( "../assets/demo", rules=rules, gen_triggers=triggers, dedupe=True, chunk=True )

3. Run the pipeline¶

Each step is wrapped in a store.create_node(...) context manager. Inside the block we write artifacts straight to the node — the store resolves paths for us via node / "file" — and attach searchable metadata with node.add_meta(...). Passing parent= links a node to the one it was derived from, so the nested loops below build the full provenance tree as a side effect of simply running the experiment.

In [3]:

scalers = {"robust": RobustScaler(), "standard": StandardScaler()}
embedders = ["umap", "tsne", "pca"]
clusterers = ["kmeans", "hdbscan"]

scale_nodes = {}
embed_nodes = {}

# Raw data ingestion
with store.create_node("raw data") as root_node:
    iris = load_iris()
    df_raw = pd.DataFrame(iris.data, columns=iris.feature_names)
    df_raw.to_csv(root_node / "data.csv", index=False)

# Two different scaling steps
for s_name, s_obj in scalers.items():
    with store.create_node("scaling", parent=root_node) as s_node:
        scaled_data = s_obj.fit_transform(df_raw)
        pd.DataFrame(scaled_data).to_csv(s_node / "scaled.csv", index=False)
        s_node.add_meta("scaler", s_name, group="Config")

    scale_nodes[s_name] = (s_node, s_obj.fit_transform(df_raw))

# Three different dimensionality reductions steps
for s_name, (s_node, scaled_data) in scale_nodes.items():
    for e_type in embedders:
        with store.create_node("embedding", parent=s_node) as e_node:
            if e_type == "pca":
                embedder = PCA(n_components=2)
            elif e_type == "tsne":
                embedder = TSNE(n_components=2, perplexity=30)
            elif e_type == "umap":
                embedder = umap.UMAP(n_components=2)

            embedding = embedder.fit_transform(scaled_data)
            np.save(e_node / "map.npy", embedding)
            e_node.add_meta("method", e_type, group="Config")
            e_node.add_meta("scaler", s_name, group="Config")

        embed_nodes[(s_name, e_type)] = (e_node, embedding)

# Two different clustering methods
for (s_name, e_type), (e_node, embedding) in embed_nodes.items():
    for c_type in clusterers:
        with store.create_node("clustering", parent=e_node) as c_node:
            if c_type == "kmeans":
                model = KMeans(n_clusters=3, n_init="auto")
            elif c_type == "hdbscan":
                model = HDBSCAN(min_cluster_size=5)

            labels = model.fit_predict(embedding)
            score = calinski_harabasz_score(embedding, labels)

            c_node.add_meta("clusterer", c_type, group="Config")
            c_node.add_meta("CH_score", round(score, 3), group="Metrics")

            fig, ax = plt.subplots(figsize=(6, 4))
            ax.scatter(embedding[:, 0], embedding[:, 1], c=labels, cmap="turbo")
            ax.set_title(f"{s_name} → {e_type} → {c_type}")
            fig.savefig(c_node / "clustering.png", bbox_inches="tight")
            plt.close(fig)

            c_node.add_meta("cluster_plot", c_node / "clustering.png", group="Results")

scalers = {"robust": RobustScaler(), "standard": StandardScaler()} embedders = ["umap", "tsne", "pca"] clusterers = ["kmeans", "hdbscan"] scale_nodes = {} embed_nodes = {} # Raw data ingestion with store.create_node("raw data") as root_node: iris = load_iris() df_raw = pd.DataFrame(iris.data, columns=iris.feature_names) df_raw.to_csv(root_node / "data.csv", index=False) # Two different scaling steps for s_name, s_obj in scalers.items(): with store.create_node("scaling", parent=root_node) as s_node: scaled_data = s_obj.fit_transform(df_raw) pd.DataFrame(scaled_data).to_csv(s_node / "scaled.csv", index=False) s_node.add_meta("scaler", s_name, group="Config") scale_nodes[s_name] = (s_node, s_obj.fit_transform(df_raw)) # Three different dimensionality reductions steps for s_name, (s_node, scaled_data) in scale_nodes.items(): for e_type in embedders: with store.create_node("embedding", parent=s_node) as e_node: if e_type == "pca": embedder = PCA(n_components=2) elif e_type == "tsne": embedder = TSNE(n_components=2, perplexity=30) elif e_type == "umap": embedder = umap.UMAP(n_components=2) embedding = embedder.fit_transform(scaled_data) np.save(e_node / "map.npy", embedding) e_node.add_meta("method", e_type, group="Config") e_node.add_meta("scaler", s_name, group="Config") embed_nodes[(s_name, e_type)] = (e_node, embedding) # Two different clustering methods for (s_name, e_type), (e_node, embedding) in embed_nodes.items(): for c_type in clusterers: with store.create_node("clustering", parent=e_node) as c_node: if c_type == "kmeans": model = KMeans(n_clusters=3, n_init="auto") elif c_type == "hdbscan": model = HDBSCAN(min_cluster_size=5) labels = model.fit_predict(embedding) score = calinski_harabasz_score(embedding, labels) c_node.add_meta("clusterer", c_type, group="Config") c_node.add_meta("CH_score", round(score, 3), group="Metrics") fig, ax = plt.subplots(figsize=(6, 4)) ax.scatter(embedding[:, 0], embedding[:, 1], c=labels, cmap="turbo") ax.set_title(f"{s_name} → {e_type} → {c_type}") fig.savefig(c_node / "clustering.png", bbox_inches="tight") plt.close(fig) c_node.add_meta("cluster_plot", c_node / "clustering.png", group="Results")

/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:722: FutureWarning: The default value of `copy` will change from False to True in 1.10. Explicitly set a value for `copy` to silence this warning.
  warn(
/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:722: FutureWarning: The default value of `copy` will change from False to True in 1.10. Explicitly set a value for `copy` to silence this warning.
  warn(
/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:722: FutureWarning: The default value of `copy` will change from False to True in 1.10. Explicitly set a value for `copy` to silence this warning.
  warn(
/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:722: FutureWarning: The default value of `copy` will change from False to True in 1.10. Explicitly set a value for `copy` to silence this warning.
  warn(
/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:722: FutureWarning: The default value of `copy` will change from False to True in 1.10. Explicitly set a value for `copy` to silence this warning.
  warn(
/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:722: FutureWarning: The default value of `copy` will change from False to True in 1.10. Explicitly set a value for `copy` to silence this warning.
  warn(

4. Visualise the lineage¶

generate_web_graph() renders the whole tree — nodes, parent/child edges, metadata and embedded figures — as a self-contained, interactive HTML page.

In [4]:

store.generate_web_graph()

store.generate_web_graph()

Out[4]:

PosixPath('../assets/demo/interactive_pipeline.html')

5. Querying the experiment¶

Because every node carries searchable metadata, we can interrogate the experiment after it has run. find_node accepts a callable for any field to express a predicate — here we pull back every clustering result whose Calinski–Harabasz score came out below 1000.

In [5]:

# Pass a lambda for richer queries than exact matching
bad_nodes = store.find_node(CH_score=lambda x: x < 1000)
bad_nodes

# Pass a lambda for richer queries than exact matching bad_nodes = store.find_node(CH_score=lambda x: x < 1000) bad_nodes

/Users/js/Documents/coding_projects/de-janked/test-venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py:3748: UserWarning: Predicate for 'CH_score' raised TypeError: '<' not supported between instances of 'NoneType' and 'int'. Node treated as non-matching.
  exec(code_obj, self.user_global_ns, self.user_ns)

Out[5]:

[Node = 2d3e1947, path = 2d3e1947, generation = 3,
 Node = c08125b5, path = c08125b5, generation = 3,
 Node = 86a52cc0, path = 86a52cc0, generation = 3,
 Node = 85769047, path = 85769047, generation = 3,
 Node = bded3ca1, path = bded3ca1, generation = 3,
 Node = 8eaaa408, path = 8eaaa408, generation = 3,
 Node = 10b06538, path = 10b06538, generation = 3,
 Node = dbbb57ad, path = dbbb57ad, generation = 3,
 Node = 8dc2e76b, path = 8dc2e76b, generation = 3]

Plain keyword arguments match exactly, and several can be combined to narrow the search. Unpacking a dict keeps things tidy when the query is assembled programmatically.

In [6]:

# Chain search terms; supply them with dictionary unpacking
search_vars = {"method": "pca", "step_type": "embedding"}

nodes_of_interest = store.find_node(**search_vars)
nodes_of_interest

# Chain search terms; supply them with dictionary unpacking search_vars = {"method": "pca", "step_type": "embedding"} nodes_of_interest = store.find_node(**search_vars) nodes_of_interest

Out[6]:

[Node = 7fbf1ed9, path = 7fbf1ed9, generation = 2,
 Node = 9da5ec69, path = 9da5ec69, generation = 2]

6. Walking a single lineage¶

get_most_recent_node() returns the last node written, and get_lineage() walks back up its ancestry to the root. Below we print the artifacts recorded at each step, from the raw data all the way down to the final clustering image.

In [ ]:

last_node = store.get_most_recent_node()
nodes_in_lineage = store.get_lineage(last_node)

for n in nodes_in_lineage:
    print(n.artifacts())

last_node = store.get_most_recent_node() nodes_in_lineage = store.get_lineage(last_node) for n in nodes_in_lineage: print(n.artifacts())

[PosixPath('/Users/js/Documents/coding_projects/de-janked/docs/assets/demo/.cache/45228-ebbd112d/80bc18e4/data.csv')]
[]
[]
[]

Reference: the rest of the store API¶

A quick tour of the other methods for when you come back to an existing store.

In [8]:

# Reopen an existing store (config is read back from disk automatically)
# store = ancestree.LineageStore("ML_pipeline")

# Search
# store.find_node(step_type="embedding")
# store.find_in_lineage(node="0d33f882", step_type="embedding")
# store.get_most_recent_node()

# Fetch by id and navigate the tree
# store.get_node(node="0d33f882")
# store.get_lineage(node="0d33f882")
# store.get_child_nodes(node="0d33f882")

# Pull a parent's artifact into the current step
# store.from_parent(node="0d33f882", filename="*.csv")

# Rebuild the index from what is on disk
# store.rebuild_db_from_disk()

# Reopen an existing store (config is read back from disk automatically) # store = ancestree.LineageStore("ML_pipeline") # Search # store.find_node(step_type="embedding") # store.find_in_lineage(node="0d33f882", step_type="embedding") # store.get_most_recent_node() # Fetch by id and navigate the tree # store.get_node(node="0d33f882") # store.get_lineage(node="0d33f882") # store.get_child_nodes(node="0d33f882") # Pull a parent's artifact into the current step # store.from_parent(node="0d33f882", filename="*.csv") # Rebuild the index from what is on disk # store.rebuild_db_from_disk()