Pipeline Configs

import laco
import laco.language as L
from omegaconf import OmegaConf
from torch import nn, optim
from torch.utils.data import DataLoader

# ============================================================
# Step 1 — A single model config
# ============================================================
model_cfg = laco.load("configs://examples/typed/mlp.py#model")
print("=== Single model ===")
# We have a recipe for one component — great.
print(list(OmegaConf.to_container(model_cfg, resolve=False).keys())[:3], "...")

# ============================================================
# Step 2 — Model + optimizer as two separate objects
# ============================================================
from laco.examples.typed.linear_regression import (
    model as lr_model_cfg,
    optimizer as lr_optim_cfg,
    OptimGroup,
)

# The two exports are *different kinds* of node:
#   - model     : a fully-formed DictConfig recipe (it has a _target_)
#   - optimizer : L.chosen(OptimGroup) — a *slot reference* into a typed
#                 group (an interpolation like '${optimgroup}'), not a
#                 standalone recipe, so it has no _target_ of its own.
print("model type        :", type(lr_model_cfg).__name__)
print("model _target_    :", lr_model_cfg._target_)   # type: ignore[union-attr]  # noqa: LACO001
print("optimizer type    :", type(lr_optim_cfg).__name__)
print("optimizer ref     :", str(lr_optim_cfg))
print("chosen variant    :", OptimGroup.sgd._target_)  # what the slot resolves to

# Problem: they are two separate objects with no shared namespace.
# A training loop would call laco.load() twice, with different fragment paths.
# There is no way to load *all components at once* or serialize the whole pipeline.

# ============================================================
# Step 3 — Bundle them into one tree
# ============================================================
# The model recipe references '${schema.*}' interpolations, and the optimizer
# is a group slot. To dump/instantiate a *self-contained* bundle, gather the
# components together with a concrete `schema` node so the references resolve
# (laco.load does this bundling for you when it reads a whole file; here we do
# it by hand for an isolated fragment).
from laco.examples.typed.linear_regression import (
    model as lr_model_cfg,
    OptimGroup,
)

train_bundle = OmegaConf.create({
    "schema": {"in_features": 4, "out_features": 1, "bias": True},
    "model": lr_model_cfg,
    "optimizer": OptimGroup.sgd,   # the concrete optimizer the slot resolves to
})

print("Bundle keys:", list(OmegaConf.to_container(train_bundle, resolve=False).keys()))
print("\n--- bundle dump (schema lets the model's interpolations resolve) ---")
print(laco.dump(train_bundle))

# Build a minimal pipeline bundle inline (no actual MNIST)
model_cfg   = L.call(nn.Linear)(in_features=784, out_features=10)
optimizer_cfg = L.partial(optim.Adam)(lr=1e-3)
loss_cfg    = L.call(nn.CrossEntropyLoss)()

pipeline = L.Dict(
    model=model_cfg,
    optimizer=optimizer_cfg,
    loss=loss_cfg,
)

print("Top-level keys:", list(OmegaConf.to_container(pipeline, resolve=False).keys()))

# Each sub-node is a fully self-contained DictConfig recipe
print("\nmodel node:")
print(laco.dump(pipeline.model))   # type: ignore[union-attr]  # noqa: LACO001

# Instantiate only what you need — no need to materialise the whole pipeline
model_obj = laco.instantiate(pipeline.model)        # type: ignore[union-attr]  # noqa: LACO001
loss_obj  = laco.instantiate(pipeline.loss)         # type: ignore[union-attr]  # noqa: LACO001

print("model:", model_obj)
print("loss: ", loss_obj)

# pipeline.py
model     = L.call(MyModel)(...)
optimizer = L.partial(optim.Adam)(lr=1e-3)
loss      = L.call(nn.CrossEntropyLoss)()

train = L.Dict(model=model, optimizer=optimizer, loss=loss)

# The import at the top of mnist_train.py:
#
#   from laco.examples.cnn_classifier import make_cnn_classifier
#
# This imports a Python *function* (not a DictConfig node).
# The function returns a DictConfig tree when called, so calling it inside
# the pipeline file is like inlining the config construction.

from laco.examples.cnn_classifier import make_cnn_classifier

# Calling the factory with explicit arguments returns a DictConfig tree
cnn_cfg = make_cnn_classifier(
    in_channels=1,
    base_channels=16,
    num_stages=2,
    num_classes=10,
)

print("type:", type(cnn_cfg))
print("_target_:", cnn_cfg._target_)  # noqa: LACO001

# How the pipeline threads interpolations through the factory:
@L.params
class pipeline_hps:
    in_channels:   int = 1
    base_channels: int = 32
    num_stages:    int = 3
    num_classes:   int = 10

# pipeline_hps.in_channels is NOT the integer 1 —
# it is the interpolation string '${pipeline_hps.in_channels}'
# (or however laco names the params namespace)
model_with_refs = make_cnn_classifier(
    in_channels=pipeline_hps.in_channels,
    base_channels=pipeline_hps.base_channels,
    num_stages=pipeline_hps.num_stages,
    num_classes=pipeline_hps.num_classes,
)

# The resulting DictConfig stores the interpolations, not literal values.
# To resolve them we bundle the params node alongside the model (calling
# pipeline_hps() materialises the namespace as a plain dict). Overriding
# pipeline_hps.num_classes=5 on this bundle would then propagate everywhere.
bundle = OmegaConf.create({"pipeline_hps": pipeline_hps(), "model": model_with_refs})
head_lines = [l for l in laco.dump(bundle).splitlines() if "num_classes" in l]
print("head out_features references:", head_lines)

# ============================================================
# Section A — Hyperparameters
# ============================================================
# @L.params creates a DictConfig namespace whose attributes are interpolation
# strings. Every other config node in the file can reference these via
# '${hps.batch_size}' etc.

from torchvision.datasets import MNIST
from torchvision import transforms

@L.params
class hps:
    data_root:     str   = "./data"
    batch_size:    int   = 64
    learning_rate: float = 1e-3
    num_workers:   int   = 0
    # Architecture hyperparameters — shared with the model factory
    in_channels:   int   = 1
    base_channels: int   = 32
    num_stages:    int   = 3
    num_classes:   int   = 10

print("hps type:", type(hps))
print("hps.batch_size (runtime):", repr(hps.batch_size))  # noqa: LACO001

# ============================================================
# Section B — Model
# ============================================================
# make_cnn_classifier is called with interpolation references.
# The factory returns a DictConfig tree where every architecture
# dimension is a '${hps.*}' string — overriding hps.num_stages=2
# at load time changes the depth of the CNN without re-running this file.

model = make_cnn_classifier(
    in_channels=hps.in_channels,
    base_channels=hps.base_channels,
    num_stages=hps.num_stages,
    num_classes=hps.num_classes,
)
print("model _target_:", model._target_)  # noqa: LACO001

# ============================================================
# Section C — Optimizer and loss
# ============================================================
# L.partial produces a "partial constructor" node (_partial_: true).
# At instantiation time you still pass model.parameters() — the optimizer
# cannot be fully instantiated without the model weights.
optimizer = L.partial(optim.Adam)(lr=hps.learning_rate)

# L.call produces a full constructor node — CrossEntropyLoss takes no
# required arguments, so it can be instantiated directly.
loss = L.call(nn.CrossEntropyLoss)()

print("optimizer _partial_:", optimizer._partial_)   # noqa: LACO001
print("loss _target_      :", loss._target_)          # noqa: LACO001

# ============================================================
# Section D — Data pipeline
# ============================================================
# L.List is a DictConfig list wrapper — it holds an ordered sequence of
# config nodes. transforms.Compose accepts a list of transforms, so we
# build the entire transform pipeline as a config tree.

_transform = L.call(transforms.Compose)(
    L.List(
        L.call(transforms.ToTensor)(),
        L.call(transforms.Normalize)(
            mean=L.List(0.1307),   # single-channel mean
            std=L.List(0.3081),    # single-channel std
        ),
    )
)

print("transform _target_:", _transform._target_)  # noqa: LACO001

# Dataset: MNIST with root, split, download, and the transform node above.
# 'download=True' is a literal bool — not an interpolation; it is part of the recipe.
dataset = L.call(MNIST)(
    root=hps.data_root,
    train=True,
    download=True,
    transform=_transform,
)

# DataLoader wraps the dataset — note that 'dataset=' receives the DictConfig
# recipe for MNIST, not an instantiated Dataset object. Hydra will
# recursively instantiate nested nodes automatically.
loader = L.call(DataLoader)(
    dataset=dataset,
    batch_size=hps.batch_size,
    shuffle=True,
    num_workers=hps.num_workers,
)

print("loader _target_  :", loader._target_)   # noqa: LACO001
print("dataset _target_ :", loader.dataset._target_)  # noqa: LACO001  # nested!

# ============================================================
# Section E — The train bundle
# ============================================================
# L.Dict collects all components under a single DictConfig root.
# This is the fragment a training script loads with:
#   cfg = laco.load("configs://examples/pipelines/mnist_train.py#train")

train = L.Dict(
    model=model,
    optimizer=optimizer,
    loss=loss,
    loader=loader,
)

print("train bundle keys:", list(OmegaConf.to_container(train, resolve=False).keys()))

# Load the actual file from the examples package (does not instantiate anything)
full_cfg = laco.load("configs://examples/pipelines/mnist_train.py")

print("Top-level fragments:", list(OmegaConf.to_container(full_cfg, resolve=False).keys()))

# Load just the train bundle fragment
train_cfg = laco.load("configs://examples/pipelines/mnist_train.py#train")
print("train bundle keys:", list(OmegaConf.to_container(train_cfg, resolve=False).keys()))
print()

# model and optimizer are immediately available as sub-configs
print("=== model ===")
print(laco.dump(train_cfg.model))   # type: ignore[union-attr]  # noqa: LACO001

# The hps namespace is also addressable as its own fragment
hps_cfg = laco.load("configs://examples/pipelines/mnist_train.py#hps")
print("=== hps ===")
print(laco.dump(hps_cfg))

@L.task
def task(
    model: nn.Module,
    optimizer: optim.Optimizer,
    loss: nn.Module,
    loader: DataLoader,
    num_steps: int = 1,
) -> None:
    ...

# Basic scalar override — change batch_size
cfg_bs128 = laco.load(
    "configs://examples/pipelines/mnist_train.py#hps",
    "hps.batch_size=128",
)
print(f"batch_size: {cfg_bs128.batch_size}")   # noqa: LACO001

# Float override — scientific notation is supported
cfg_lr = laco.load(
    "configs://examples/pipelines/mnist_train.py#hps",
    "hps.learning_rate=5e-4",
)
print(f"learning_rate: {cfg_lr.learning_rate}")  # noqa: LACO001

# Multiple overrides in a single call
cfg_multi = laco.load(
    "configs://examples/pipelines/mnist_train.py#hps",
    "hps.batch_size=256",
    "hps.num_stages=2",
    "hps.base_channels=16",
)
print(f"batch_size={cfg_multi.batch_size}  "
      f"num_stages={cfg_multi.num_stages}  "
      f"base_channels={cfg_multi.base_channels}")  # noqa: LACO001

# Show what changed between base and overridden configs
base_hps = laco.load("configs://examples/pipelines/mnist_train.py#hps")
ovrd_hps = laco.load(
    "configs://examples/pipelines/mnist_train.py#hps",
    "hps.batch_size=128",
    "hps.learning_rate=5e-4",
)

for key in ("batch_size", "learning_rate", "num_stages"):
    before = OmegaConf.select(base_hps, key)
    after  = OmegaConf.select(ovrd_hps, key)
    changed = "<-- changed" if before != after else ""
    print(f"  {key:<18} {before!r:>12}  →  {after!r:<12} {changed}")

# URL-style: ?key=value before the # fragment separator
cfg_url = laco.load(
    "configs://examples/pipelines/mnist_train.py"
    "?hps.batch_size=32&hps.learning_rate=2e-3"
    "#hps"
)
print(f"batch_size={cfg_url.batch_size}  learning_rate={cfg_url.learning_rate}")  # noqa: LACO001

python -m laco.examples.pipelines.mnist_train num_steps=2

# From the bottom of mnist_train.py:
#
#   if __name__ == "__main__":
#       import laco
#       cfg = laco.load(__file__ + "#train")
#       task(cfg)
#
# Why __file__ instead of "configs://examples/pipelines/mnist_train.py"?
#
# __file__ is the absolute path on disk.  laco.load() accepts both:
#   - a "configs://" URL (resolved against the package search path)
#   - a raw filesystem path (resolved directly)
# Using __file__ makes the script relocatable — it works even if the module
# is installed in a virtualenv, editable-installed, or symlinked.

import laco.examples.pipelines.mnist_train as _mt_module
print("Module file:", _mt_module.__file__)

# Load via filesystem path (same as __file__ + "#train" from inside the module)
import importlib.resources, pathlib

module_path = pathlib.Path(_mt_module.__file__)
cfg_via_file = laco.load(str(module_path) + "#hps")
print("Loaded via __file__ path — batch_size:", cfg_via_file.batch_size)  # noqa: LACO001

# Load the train bundle
train_cfg = laco.load(
    "configs://examples/pipelines/mnist_train.py#train",
    "hps.base_channels=16",   # smaller model for faster instantiation in this notebook
    "hps.num_stages=2",
)

# Instantiate each component individually
model_obj     = laco.instantiate(train_cfg.model)      # type: ignore[union-attr]  # noqa: LACO001
optimizer_fn  = laco.instantiate(train_cfg.optimizer)  # type: ignore[union-attr]  # noqa: LACO001
loss_obj      = laco.instantiate(train_cfg.loss)       # type: ignore[union-attr]  # noqa: LACO001

# optimizer is a partial — call it with model.parameters()
opt = optimizer_fn(model_obj.parameters())

print("model     :", type(model_obj).__name__)
print("optimizer :", type(opt).__name__)
print("loss      :", type(loss_obj).__name__)

# Verify model works with a dummy forward pass
import torch
dummy = torch.randn(2, 1, 28, 28)  # batch=2, channels=1, 28x28
out = model_obj(dummy)
print("output shape:", out.shape)   # should be [2, 10]

# Inspect the loader config — it contains a nested dataset node
loader_cfg = laco.load("configs://examples/pipelines/mnist_train.py#loader")
print("=== loader config ===")
print(laco.dump(loader_cfg))

# The nesting depth: loader → dataset → transform → [ToTensor, Normalize]
# laco.instantiate(loader_cfg) would:
#   1. Build Normalize(mean=[0.1307], std=[0.3081])
#   2. Build ToTensor()
#   3. Build Compose([ToTensor(), Normalize(...)])
#   4. Build MNIST(root=..., transform=Compose(...))
#   5. Build DataLoader(dataset=MNIST(...), batch_size=64, ...)
#
# This cell is *illustrative* — MNIST download would be triggered.
# Uncomment the line below in an environment with internet access:

# loader_obj = laco.instantiate(loader_cfg)
# print("loader:", loader_obj)

print("(Skipped — would download MNIST.  Run in a connected environment to verify.)")