Notebook

Tasks and the App Loop

Series: laco tutorial notebooks
Prerequisites: 01.why-laco.ipynb through 08.pipeline-configs.ipynb (especially 03.lazy-call-and-partial.ipynb, 04.hyperparameters-and-interpolation.ipynb, 07.typed-groups-and-schemas.ipynb)
Dependencies: torch, torchvision


In previous notebooks you learned how to build configs: L.call, L.partial, @L.params, groups, defaults. This notebook covers the other half: how to run them.

The key concept is the task function, a callable that accepts a DictConfig and automatically receives its fields as fully-instantiated Python objects. Two decorators power this:

DecoratorWhat it does
@L.taskUnwraps a DictConfig into typed kwargs for one function
@laco.main(config_name=...)Hydra app entry point; composes config from file + CLI, then calls @L.task

By the end you will understand exactly what happens between laco.load("train.py") and model.train().

import inspect
import functools

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

Section 1: The problem @L.task solves

Suppose you have a training config loaded via laco.load("train.py") that contains model, optimizer, loss, and num_steps fields. Without @L.task, wiring them to a function is repetitive boilerplate:

# WITHOUT @L.task — verbose, error-prone, and not refactor-safe
def train_raw(cfg):
    model          = laco.instantiate(cfg.model)
    optimizer_fact = laco.instantiate(cfg.optimizer)   # functools.partial
    optimizer      = optimizer_fact(model.parameters())
    loss_fn        = laco.instantiate(cfg.loss)
    num_steps      = laco.instantiate(cfg.num_steps)   # plain int — also goes through instantiate

    model.train()
    # ... training loop using model, optimizer, loss_fn, num_steps
    print("train_raw: instantiation done")

print(inspect.getsource(train_raw))
Output
def train_raw(cfg):
    model          = laco.instantiate(cfg.model)
    optimizer_fact = laco.instantiate(cfg.optimizer)   # functools.partial
    optimizer      = optimizer_fact(model.parameters())
    loss_fn        = laco.instantiate(cfg.loss)
    num_steps      = laco.instantiate(cfg.num_steps)   # plain int — also goes through instantiate

    model.train()
    # ... training loop using model, optimizer, loss_fn, num_steps
    print("train_raw: instantiation done")

Pain points:

  • Every field must be manually laco.instantiated.
  • If you rename a config key, you must update both the config file and every cfg.field_name access.
  • There is no static type information on cfg.model, so the IDE can't help.

With @L.task, laco reads your function's signature and does the wiring for you:

# WITH @L.task — signature-driven, type-annotated, refactor-safe
@L.task
def train(model: nn.Module, optimizer, loss: nn.Module, num_steps: int = 1):
    # All fields already instantiated by @L.task!
    # optimizer is a functools.partial — call it with model.parameters()
    opt = optimizer(model.parameters())
    model.train()
    print(f"train: model={type(model).__name__}, loss={type(loss).__name__}, steps={num_steps}")

# The decorated function now accepts a DictConfig, not positional args:
print("Signature of wrapped 'train':", inspect.signature(train))
print("Has _laco_task marker:", getattr(train, '_laco_task', False))
Output
Signature of wrapped 'train': (model: torch.nn.modules.module.Module, optimizer, loss: torch.nn.modules.module.Module, num_steps: int = 1)
Has _laco_task marker: True

Section 2: @L.task internals, step by step

When you write @L.task on a function, laco does the following at decoration time (once, not on every call):

  1. inspect.signature(train) → discovers parameter names: model, optimizer, loss, num_steps

Then at call time (each time you call train(cfg)):

  1. For each parameter name p:
    • OmegaConf.select(cfg, p) → retrieves the sub-tree (or SENTINEL if absent)
    • laco.instantiate(sub_tree) → materializes it: DictConfig{_target_: nn.Linear ...}nn.Linear object
  2. If a required parameter is missing from the config and has no default: TypeError with clear message
  3. train(model=model_obj, optimizer=partial_obj, loss=loss_obj, num_steps=1): the original body runs

VAR_POSITIONAL (*args) and VAR_KEYWORD (**kwargs) parameters are skipped: @L.task only maps named parameters.

# Manually reproduce what @L.task does, for pedagogical clarity:
from omegaconf import OmegaConf

# Build a toy config manually (no file needed)
cfg_manual = OmegaConf.create({
    "value": 42,
    "scale": 2.5,
})

@L.task
def process(value: int, scale: float = 1.0) -> float:
    result = value * scale
    print(f"process(value={value}, scale={scale}) = {result}")
    return result

# The task reads 'value' and 'scale' from the DictConfig:
result = process(cfg_manual)
print(f"Result: {result}")
Output
process(value=42, scale=2.5) = 105.0
Result: 105.0
# What happens when a required parameter is missing?
cfg_incomplete = OmegaConf.create({"scale": 3.0})  # 'value' is absent and has no default

try:
    process(cfg_incomplete)
except TypeError as e:
    print("TypeError caught (required param missing):")
    print(e)
Output
TypeError caught (required param missing):
@L.task(process): required parameter 'value' not found in config. Available keys: ['scale']
# Parameters with defaults: omitting 'scale' is fine — the function default applies.
cfg_no_scale = OmegaConf.create({"value": 10})
result_default = process(cfg_no_scale)
print(f"Result with default scale: {result_default}")
Output
process(value=10, scale=1.0) = 10.0
Result with default scale: 10.0

Section 3: The optimizer partial pattern

The most important @L.task use case is the optimizer partial pattern. Optimizers like Adam need model.parameters(), which only exists after the model is built. laco solves this with L.partial: the config stores _partial_: true, laco.instantiate returns a functools.partial, and the task body calls it with model.parameters().

This is exactly how the MNIST pipeline works:

# Build a small config that mimics the MNIST pipeline structure:
model_cfg     = L.call(nn.Linear)(in_features=4, out_features=2)
optimizer_cfg = L.partial(optim.SGD)(lr=1e-2, momentum=0.9)
loss_cfg      = L.call(nn.CrossEntropyLoss)()

combined_cfg = OmegaConf.create({
    "model":     OmegaConf.to_container(model_cfg,     resolve=False),
    "optimizer": OmegaConf.to_container(optimizer_cfg, resolve=False),
    "loss":      OmegaConf.to_container(loss_cfg,      resolve=False),
    "num_steps": 2,
})

@L.task
def mini_train(model: nn.Module, optimizer, loss: nn.Module, num_steps: int = 1):
    # At this point:
    #   model     — a real nn.Linear
    #   optimizer — a functools.partial(SGD, lr=0.01, momentum=0.9)
    #   loss      — a real nn.CrossEntropyLoss
    print(f"model type     : {type(model).__name__}")
    print(f"optimizer type : {type(optimizer).__name__}")
    print(f"loss type      : {type(loss).__name__}")
    print(f"num_steps      : {num_steps}")

    # Now create the full optimizer using model.parameters():
    opt = optimizer(model.parameters())  # functools.partial called here
    print(f"opt type       : {type(opt).__name__}")
    print(f"opt lr         : {opt.param_groups[0]['lr']}")

mini_train(combined_cfg)
Output
model type     : Linear
optimizer type : partial
loss type      : CrossEntropyLoss
num_steps      : 2
opt type       : SGD
opt lr         : 0.01

Key insight: The optimizer parameter in the function body is a functools.partial, not yet a full optimizer. The task body must call optimizer(model.parameters()) to create the optimizer. This two-step pattern is intentional: it keeps the config pure (no live objects) while still allowing the optimizer to capture the model's parameters at the right time.

Section 4: @laco.main, the full Hydra app

@laco.main is the entry point for scripts you run from the command line. It wraps @hydra.main and automatically applies @L.task semantics. The result: Hydra handles config composition and CLI overrides; laco handles instantiation and kwarg mapping.

Typical usage in a training script:

# Illustrative — the pattern used in mnist_train.py and clm_finetune.py:

# import laco
# import laco.language as L
# from torch import nn, optim
#
# @laco.main(config_name="train", config_path="configs")
# @L.task
# def run(model: nn.Module, optimizer, loss: nn.Module, num_steps: int = 10):
#     opt = optimizer(model.parameters())
#     model.train()
#     for step in range(num_steps):
#         ...  # training step
#
# if __name__ == "__main__":
#     run()  # Hydra takes over; parses argv; composes config; calls @L.task wrapper

print("Pattern shown above (not executed in notebook — requires __main__ guard)")
print()
print("Execution flow:")
print("  1. run()             → Hydra parses sys.argv")
print("  2. Hydra composes    → builds DictConfig from config_name + overrides")
print("  3. @L.task unwraps   → instantiates each field matching a parameter")
print("  4. run body executes → receives typed, instantiated objects")
Output
Pattern shown above (not executed in notebook — requires __main__ guard)

Execution flow:
  1. run()             → Hydra parses sys.argv
  2. Hydra composes    → builds DictConfig from config_name + overrides
  3. @L.task unwraps   → instantiates each field matching a parameter
  4. run body executes → receives typed, instantiated objects

Implicit @L.task

If the function is not already wrapped with @L.task, @laco.main applies it implicitly. Both forms below are equivalent:

# Form 1: explicit @L.task + @laco.main
# @laco.main(config_name="train")
# @L.task
# def run_explicit(model: nn.Module, num_steps: int = 10): ...

# Form 2: only @laco.main — @L.task is applied automatically
# @laco.main(config_name="train")
# def run_implicit(model: nn.Module, num_steps: int = 10): ...

# The check inside laco.main:
# task_fn = func if getattr(func, '_laco_task', False) else _task(func)

# Demonstrate the check:
@L.task
def already_a_task(x: int): pass

def not_yet_a_task(x: int): pass

print("already_a_task._laco_task :", getattr(already_a_task, '_laco_task', False))
print("not_yet_a_task._laco_task :", getattr(not_yet_a_task,  '_laco_task', False))
Output
already_a_task._laco_task : True
not_yet_a_task._laco_task : False

@laco.main parameters

ParameterDefaultForwarded to
config_name"config"@hydra.main(config_name=...)
config_pathNone@hydra.main(config_path=...)
version_baseNone@hydra.main(version_base=...)

When config_path=None, Hydra uses its own search path (config files next to the script, or on HYDRA_CONFIG_PATH). For laco config files, the configs:// scheme (from laco.handler) provides an alternative path resolver.

Section 5: The MNIST training pipeline

The laco.examples.pipelines.mnist_train module is the canonical end-to-end example. Let's read its source to see @L.task in a real pipeline:

import inspect
import laco.examples.pipelines.mnist_train as mnist

print("=== mnist_train.task (the @L.task entry point) ===")
print(inspect.getsource(mnist.task))
Output
=== mnist_train.task (the @L.task entry point) ===
@L.task
def task(
    model: nn.Module,
    optimizer: optim.Optimizer,
    loss: nn.Module,
    loader: DataLoader,
    num_steps: int = 1,
) -> None:
    """Run ``num_steps`` training iterations (default 1 for smoke testing).

    Invoke via::

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

    opt = optimizer(model.parameters())  # type: ignore[operator]
    model.train()
    loader_iter = iter(loader)
    for step in range(num_steps):
        try:
            images, labels = next(loader_iter)
        except StopIteration:
            loader_iter = iter(loader)
            images, labels = next(loader_iter)
        opt.zero_grad()
        out = model(images)
        loss_val = loss(out, labels)
        loss_val.backward()
        opt.step()
        _ = step  # suppress unused-variable lint

# Load the training sub-config (model + optimizer + loss + loader)
cfg_train = laco.load("configs://examples/pipelines/mnist_train.py#train")

print("Keys in the 'train' sub-config:")
for k in cfg_train.keys():
    print(f"  {k}")
Output
Keys in the 'train' sub-config:
  _target_
  _convert_
  model
  optimizer
  loss
  loader
# Instantiate only the model (cheap — no MNIST download needed)
model = laco.instantiate(cfg_train.model)
print("Model type:", type(model).__name__)
print("Model:", model)
Output
Model type: Sequential
Model: Sequential(
  (stem): Sequential(
    (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
    (2): ReLU(inplace=True)
  )
  (stages): Sequential(
    (0): Sequential(
      (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (1): Sequential(
      (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (2): Sequential(
      (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
  )
  (head): Sequential(
    (0): AdaptiveAvgPool2d(output_size=1)
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=32, out_features=10, bias=True)
  )
)
# Run the pipeline via subprocess (shows Hydra + @L.task wiring end-to-end)
# We use num_steps=1 for a fast smoke test (downloads MNIST if not cached)
import subprocess
result = subprocess.run(
    ["python", "-m", "laco.examples.pipelines.mnist_train", "num_steps=1"],
    capture_output=True, text=True,
    cwd="/home/khwstolle/Projects/research/laco",
    timeout=120,
)
stdout_tail = result.stdout[-600:] if result.stdout else ""
stderr_tail = result.stderr[-600:] if result.stderr else ""
print("--- stdout ---")
print(stdout_tail or "(empty)")
print("--- stderr ---")
print(stderr_tail or "(empty)")
print("Return code:", result.returncode)
Output
--- stdout ---
(empty)
--- stderr ---
(empty)
Return code: 0

Section 6: @L.task data flow

The complete data flow from a DictConfig to a running training function.

StepInputOperationOutput
1 (at decoration)train's signatureinspect.signature(train)Named parameters: model, optimizer, loss, num_steps
2 (per call, per param)DictConfig + param nameOmegaConf.select(cfg, name)Sub-tree for that field
3 (per call, per param)Sub-treelaco.instantiate(sub_tree)Instantiated object (e.g. modelnn.Module, optimizerfunctools.partial)
4 (per call)All instantiated kwargstrain(model=..., optimizer=..., loss=..., num_steps=...)The original function body runs

optimizer is a functools.partial, not a live optimizer — the function body still calls optimizer(model.parameters()) to get the real object.

Section 7: Contrast with raw @hydra.main

The value of @laco.main + @L.task becomes clear when you compare it with plain @hydra.main:

# ============================================================
# WITHOUT laco: raw @hydra.main — ~20 lines of boilerplate
# ============================================================

# from hydra.utils import instantiate
# import hydra
# from omegaconf import DictConfig
#
# @hydra.main(config_name="train", config_path="configs", version_base=None)
# def run_hydra(cfg: DictConfig) -> None:
#     # Every field requires a manual instantiate call
#     model     = instantiate(cfg.model)
#     opt_cfg   = instantiate(cfg.optimizer)          # functools.partial
#     optimizer = opt_cfg(model.parameters())         # finish construction
#     loss      = instantiate(cfg.loss)
#     loader    = instantiate(cfg.dataset)            # also instantiate dataset
#     loader    = instantiate(cfg.loader,             # and loader separately
#                             dataset=dataset)
#     num_steps = cfg.hps.num_steps                   # no instantiate for primitives
#
#     model.train()
#     for step in range(num_steps):
#         images, labels = next(iter(loader))
#         loss_val = loss(model(images), labels)
#         # ...
#
# if __name__ == "__main__":
#     run_hydra()

raw_lines = """
@hydra.main(config_name="train", config_path="configs", version_base=None)
def run(cfg: DictConfig) -> None:
    model     = instantiate(cfg.model)
    opt_cfg   = instantiate(cfg.optimizer)
    optimizer = opt_cfg(model.parameters())
    loss      = instantiate(cfg.loss)
    loader    = instantiate(cfg.loader)
    num_steps = cfg.hps.num_steps
    # ... actual training code ...
"""

# ============================================================
# WITH laco: @laco.main + @L.task — 5 lines
# ============================================================

laco_lines = """
@laco.main(config_name="train")
@L.task
def run(model: nn.Module, optimizer, loss: nn.Module,
        loader: DataLoader, num_steps: int = 10):
    opt = optimizer(model.parameters())  # optimizer is a partial
    # ... actual training code ...
"""

print("=== Raw @hydra.main ===")
print(raw_lines)
print("=== @laco.main + @L.task ===")
print(laco_lines)
Output
=== Raw @hydra.main ===

@hydra.main(config_name="train", config_path="configs", version_base=None)
def run(cfg: DictConfig) -> None:
    model     = instantiate(cfg.model)
    opt_cfg   = instantiate(cfg.optimizer)
    optimizer = opt_cfg(model.parameters())
    loss      = instantiate(cfg.loss)
    loader    = instantiate(cfg.loader)
    num_steps = cfg.hps.num_steps
    # ... actual training code ...

=== @laco.main + @L.task ===

@laco.main(config_name="train")
@L.task
def run(model: nn.Module, optimizer, loss: nn.Module,
        loader: DataLoader, num_steps: int = 10):
    opt = optimizer(model.parameters())  # optimizer is a partial
    # ... actual training code ...

What laco saves you:

  • Every instantiate(cfg.field_name) call is gone: laco generates them from the signature.
  • The function parameters are type-annotated, so pyright knows model is nn.Module.
  • If you add, remove, or rename a config field, only the function signature changes, not a scattered list of cfg.x accesses.
  • The raw @hydra.main approach also requires knowing which fields need instantiate vs direct access (cfg.hps.num_steps vs instantiate(cfg.model)). @L.task uniformly applies laco.instantiate to everything (which is a no-op for plain primitives).

Section 8: Multirun and sweeps

@laco.main delegates directly to Hydra, so all of Hydra's multirun, sweeper, and launcher plugins work out of the box. A single CLI invocation can launch a grid search across optimizers and learning rates:

# Illustrative CLI forms — not executed in this notebook.
# They require a running Python environment with the config files available.

multirun_examples = """
# === Hydra multirun via @laco.main ===

# Grid search: 2 optimizers × 2 learning rates = 4 runs
python train.py -m \
    optimizer=sgd,adam \
    hps.learning_rate=1e-2,1e-3
# Produces:
#   run 1: optimizer=sgd,  lr=1e-2
#   run 2: optimizer=sgd,  lr=1e-3
#   run 3: optimizer=adam, lr=1e-2
#   run 4: optimizer=adam, lr=1e-3

# Optuna sweep (requires hydra-optuna-sweeper plugin):
python train.py -m \
    hydra/sweeper=optuna \
    'hps.learning_rate=interval(1e-4, 1e-1)' \
    hydra.sweeper.n_trials=20

# SLURM cluster (requires hydra-submitit-launcher plugin):
python train.py -m \
    hydra/launcher=submitit_slurm \
    optimizer=sgd,adam \
    hps.learning_rate=1e-3,1e-4
"""

print(multirun_examples)
Output

# === Hydra multirun via @laco.main ===

# Grid search: 2 optimizers × 2 learning rates = 4 runs
python train.py -m     optimizer=sgd,adam     hps.learning_rate=1e-2,1e-3
# Produces:
#   run 1: optimizer=sgd,  lr=1e-2
#   run 2: optimizer=sgd,  lr=1e-3
#   run 3: optimizer=adam, lr=1e-2
#   run 4: optimizer=adam, lr=1e-3

# Optuna sweep (requires hydra-optuna-sweeper plugin):
python train.py -m     hydra/sweeper=optuna     'hps.learning_rate=interval(1e-4, 1e-1)'     hydra.sweeper.n_trials=20

# SLURM cluster (requires hydra-submitit-launcher plugin):
python train.py -m     hydra/launcher=submitit_slurm     optimizer=sgd,adam     hps.learning_rate=1e-3,1e-4

The -m flag activates Hydra's multirun mode. In this mode, Hydra expands the comma-separated values into a Cartesian product of runs (or hands them off to a configured sweeper). Each run calls @L.task with its own config. The training script itself needs no changes.

This is a major advantage of building on Hydra: the same script works for single-run development and large-scale hyperparameter searches without modification.

Recap

ConceptKey fact
@L.taskReads inspect.signature; calls OmegaConf.select + laco.instantiate for each named param
wrapper._laco_task = TrueMarker used by @laco.main to avoid double-wrapping
optimizer param typefunctools.partial; must call optimizer(model.parameters()) inside the body
@laco.mainWraps @hydra.main; applies @L.task implicitly if needed
MultirunPass -m on CLI; Hydra sweeps, @laco.main is transparent
Missing required paramTypeError with clear message listing available keys

Next: 10.tracing.ipynb covers the tracing layer: @L.configurable and L.trace, which let you write normal Python constructors and automatically capture them as config nodes.