Production Patterns

import laco
import laco.language as L
import torch.nn as nn

import pathlib
import tempfile

# Load the MNIST training config (fast — no MNIST download).
#
# We load the WHOLE config file (every `__all__` export) rather than just the
# `#train` subtree. The components in `train` interpolate hyperparameters via
# `${hps.batch_size}` etc., and those `${hps.*}` references only resolve when the
# `hps` node is a SIBLING in the same config tree. `laco.load(file)` bundles all
# exports into one root tree, so `hps` travels alongside `train` and the saved
# YAML is fully self-contained. (Saving the bare `#train` fragment would strip
# `hps`, leaving dangling interpolations that break on reload.)
cfg = laco.load("configs://examples/pipelines/mnist_train.py")

# Create a run directory and save the config
run_dir = pathlib.Path(tempfile.mkdtemp())
saved   = laco.save(cfg, run_dir / "config.yaml")
print(f"Config saved to: {run_dir / 'config.yaml'}")
print()

# Peek at the saved YAML
yaml_content = (run_dir / "config.yaml").read_text()
print("=== Saved config.yaml (first 60 lines) ===")
for line in yaml_content.splitlines()[:60]:
    print(line)

# Reproduce: load from the saved YAML — no original .py file needed
cfg_repro = laco.load(run_dir / "config.yaml")

# Verify the round-trip is perfect
assert laco.dump(cfg) == laco.dump(cfg_repro), "Round-trip mismatch!"
print("Config round-trip verified.")
print()
print("The saved YAML contains all _target_ strings needed to reproduce the experiment.")
print("Share config.yaml alongside your model checkpoint for full reproducibility.")

# Cleanup
import shutil
shutil.rmtree(run_dir)

@laco.main(config_name="train")
@L.task
def run(model: nn.Module, optimizer, loss: nn.Module,
        loader, num_steps: int = 10, output_dir: str = "outputs") -> None:
    run_dir = pathlib.Path(output_dir)
    run_dir.mkdir(parents=True, exist_ok=True)

    # Save BEFORE any training — even a crash partway through preserves the config
    laco.save(cfg, run_dir / "config.yaml")  # cfg available as closure from @L.task

    opt = optimizer(model.parameters())
    # ... training ...

import os

# ---- With LACO_STRICT_NODES=1 ----
os.environ["LACO_STRICT_NODES"] = "1"

# Re-import to pick up the env var (laco._strict checks it at module load time)
import importlib
import laco._strict
importlib.reload(laco._strict)
import laco.language as L_strict
importlib.reload(L_strict)

cfg_strict = L_strict.call(nn.Linear)(in_features=16, out_features=8)
print("Type of cfg_strict:", type(cfg_strict).__name__)

# Key access is allowed (normal dict-key semantics):
print("in_features (key access):", cfg_strict["in_features"])

# Attribute access RAISES with a clear message:
try:
    _ = cfg_strict.in_features   # noqa: LACO001
except AttributeError as e:
    print("\nAttributeError (LACO_STRICT_NODES):")
    print(str(e)[:200])

# Disable LACO_STRICT_NODES for the rest of the notebook
os.environ.pop("LACO_STRICT_NODES", None)
importlib.reload(laco._strict)
importlib.reload(L_strict)
print("LACO_STRICT_NODES disabled — normal behaviour restored.")

# What LACO001 catches:

example_bad = '''
import laco.language as L
import torch.nn as nn

cfg = L.call(nn.Linear)(in_features=16, out_features=8)
# BAD: cfg is a DictConfig, but the IDE (and this code) treats it as nn.Linear
print(cfg.in_features)  # LACO001: attribute access on a lie-typed node
'''

example_good = '''
import laco.language as L
import torch.nn as nn

cfg = L.call(nn.Linear)(in_features=16, out_features=8)
# GOOD 1: suppressed with noqa comment (intentional config key access)
print(cfg.in_features)  # noqa: LACO001 -- intentional: exercising real key access

# GOOD 2: use dict-key syntax (always safe, no suppression needed)
print(cfg["in_features"])  # plain dict key — no LACO001
'''

print("=== Bad pattern (triggers LACO001) ===")
print(example_bad)
print("=== Good patterns ===")
print(example_good)

import subprocess

# Run laco-lint on the examples directory
result = subprocess.run(
    ["laco-lint", "sources/laco/examples/"],
    capture_output=True, text=True,
    cwd="/home/khwstolle/Projects/research/laco",
    timeout=30,
)
print("laco-lint stdout:")
print(result.stdout or "(no violations found)")
if result.stderr:
    print("laco-lint stderr:")
    print(result.stderr[:300])
print("Return code:", result.returncode, "(0 = clean, nonzero = violations found)")

# .github/workflows/lint.yml
- name: laco-lint
  run: laco-lint sources/ tests/

import laco._lazy
import warnings

# Save the original (identity) hook so we can restore it later
original_migrate = laco._lazy._DEFAULT_MIGRATE_TARGET

# Define a migration table for a hypothetical 1.0 → 2.0 rename:
MIGRATIONS = {
    "mypackage.old.ResNet": "mypackage.models.ResNet",
    "mypackage.utils.OldNorm": "mypackage.layers.LayerNorm",
}

def my_migration_hook(target):
    if target in MIGRATIONS:
        new_target = MIGRATIONS[target]
        warnings.warn(
            f"'{target}' is deprecated; automatically remapped to '{new_target}'. "
            f"Update your YAML configs to use the new path.",
            DeprecationWarning,
            stacklevel=4,
        )
        return new_target
    return target

# Install the hook:
laco._lazy.migrate_target = my_migration_hook

# Test the hook:
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    result = laco._lazy.migrate_target("mypackage.old.ResNet")
    unchanged = laco._lazy.migrate_target("mypackage.models.NewNet")

print(f"Old path remapped to: {result}")
print(f"Unknown path unchanged: {unchanged}")
if w:
    print(f"Warning: {w[0].message}")

# Restore the identity hook:
laco._lazy.migrate_target = original_migrate
print("migrate_target restored to identity hook.")

import laco.compat
import warnings

# The hook is installed on import.
# It does NOT raise an error for the old format — it just warns:
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    result = laco._lazy.migrate_target("laco.ops.partial")

print(f"Target after compat hook: {result!r}")  # 'laco.ops.partial' — returned unchanged
if w:
    print(f"Warning category: {w[0].category.__name__}")
    print(f"Warning message: {w[0].message}")

# Uninstall when done (restores identity hook):
laco.compat.uninstall()
print("compat hook uninstalled.")
print("migrate_target is now identity:",
      laco._lazy.migrate_target is laco._lazy._DEFAULT_MIGRATE_TARGET)

# Illustrate the bind_strict + use API:
from laco.language import bind_strict, use, bind, Defaults, self_, slot

# Define two groups with the same element type:
class ActivationGroup(L.Group[nn.Module]):
    relu    = L.call(nn.ReLU)()
    gelu    = L.call(nn.GELU)()
    silu    = L.call(nn.SiLU)()

class NormGroup(L.Group[nn.Module]):
    layer   = L.call(nn.LayerNorm)(normalized_shape=256)

@L.config
class BackboneSchema:
    hidden:     int       = 256
    activation: nn.Module = slot(ActivationGroup)

# L.bind — runtime check only:
binding_ok   = bind(BackboneSchema.activation, ActivationGroup.gelu)
print("bind OK  :", binding_ok)

# L.bind with mismatched groups raises at config-build time:
try:
    bind(BackboneSchema.activation, NormGroup.layer)  # wrong group!
except TypeError as e:
    print("bind mismatch TypeError:", str(e)[:120])

# L.bind_strict — also requires L.use() wrapper on the entry:
binding_strict = bind_strict(BackboneSchema.activation, use(ActivationGroup.relu))
print("bind_strict OK:", binding_strict)

# At STATIC analysis time (pyright), this would also be caught:
# bind_strict(BackboneSchema.activation, use(NormGroup.layer))  # <-- pyright error
# The invariant EntryRef[T] parameter forces the type mismatch to surface.
print()
print("L.use() at runtime is identity (returns the entry unchanged):")
import laco._groups
entry = ActivationGroup.relu
wrapped = use(entry)
print(f"  entry is use(entry): {entry is wrapped}")

# Base group (defined in core library):
class LossGroup(L.Group[nn.Module]):
    cross_entropy  = L.call(nn.CrossEntropyLoss)()
    mse            = L.call(nn.MSELoss)()

# Extension group (defined in a downstream package or plugin):
class MoreLosses(LossGroup, mode="extend"):
    huber          = L.call(nn.HuberLoss)()
    bce_logits     = L.call(nn.BCEWithLogitsLoss)()

# All entries are accessible from the extension:
print("LossGroup entries    :", [k for k in vars(LossGroup) if not k.startswith('_')])
print("MoreLosses entries   :", [k for k in vars(MoreLosses) if not k.startswith('_')])
print()
# Can build a node for the new entries:
print(laco.dump(MoreLosses.huber))

from laco.language import delete, Defaults, self_

# L.delete produces the Hydra ~group removal entry:
del_entry = delete("logger")
print("L.delete('logger'):", del_entry)
print("Type:", type(del_entry).__name__)

# Using it in a Defaults list:
# defaults = Defaults(self_, delete("logger"), bind(Schema.optimizer, OptimizerGroup.adam))
# This would:
#   1. position self_ (the schema config) in the merge order
#   2. remove the 'logger' group from the composed config
#   3. select 'adam' from OptimizerGroup
print()
print("Typical Defaults usage:")
print("  Defaults(self_,")
print("          delete('logger'),")
print("          bind(Schema.optimizer, OptimizerGroup.adam))")

import torch.nn as nn
import torch.optim as optim
import laco
import laco.language as L
from laco.language import bind, Defaults, self_, slot

# ================================================================
# 1. Config groups: swappable components
# ================================================================

class OptimizerGroup(L.Group[object]):  # object base since partial is not nn.Module
    adam = L.partial(optim.Adam)(lr=1e-3, weight_decay=1e-5)
    sgd  = L.partial(optim.SGD)(lr=1e-2, momentum=0.9)

class ModelGroup(L.Group[nn.Module]):
    small  = L.call(nn.Linear)(in_features=64,  out_features=10)
    medium = L.call(nn.Linear)(in_features=256, out_features=10)

# ================================================================
# 2. Typed schema: validates fields at compose time
# ================================================================

@L.config
class TrainSchema:
    epochs:     int    = 10
    batch_size: int    = 32
    optimizer:  object = slot(OptimizerGroup)
    model:      nn.Module = slot(ModelGroup)

# ================================================================
# 3. Defaults list: select concrete entries
# ================================================================

defaults = Defaults(
    self_,
    bind(TrainSchema.optimizer, OptimizerGroup.adam),
    bind(TrainSchema.model,     ModelGroup.small),
)

print("Schema:", TrainSchema)
print("Defaults:", defaults)

# ================================================================
# 4. Task function: receives instantiated objects
# ================================================================

@L.task
def run(model: nn.Module, optimizer, epochs: int = 10, batch_size: int = 32):
    """Training task — parameters automatically instantiated by @L.task."""
    opt = optimizer(model.parameters())
    print(f"  model     : {model}")
    print(f"  optimizer : {type(opt).__name__}  lr={opt.param_groups[0]['lr']}")
    print(f"  epochs    : {epochs}")
    print(f"  batch_size: {batch_size}")
    # ... actual training loop here ...

# ================================================================
# 5. Build and run
# ================================================================

# Build a DictConfig manually for demonstration (in production: @laco.main does this)
from omegaconf import OmegaConf
cfg_run = OmegaConf.create({
    "model":      OmegaConf.to_container(ModelGroup.small,     resolve=False),
    "optimizer":  OmegaConf.to_container(OptimizerGroup.adam,  resolve=False),
    "epochs":     5,
    "batch_size": 64,
})

print("=== Running task ===")
run(cfg_run)

# ================================================================
# 6. Save for reproducibility
# ================================================================

import pathlib, tempfile

run_dir = pathlib.Path(tempfile.mkdtemp())
laco.save(cfg_run, run_dir / "config.yaml")

print("Saved config:")
print((run_dir / "config.yaml").read_text())

# Later: reload and run again
cfg_repro = laco.load(run_dir / "config.yaml")
print("=== Reproduced run ===")
run(cfg_repro)

import shutil
shutil.rmtree(run_dir)