[doc] pytorch lightning example

.github/workflows/test-type-lint.yaml

@@ -48,6 +48,7 @@ jobs:
       run: |
         source activate ./ci_env
         pip install -e .[dev]
+        pip install scikit-learn lightning torchvision  # for docs
 
     - name: Print installed packages
       run: |

docs/infra/example_lightning.py

@@ -0,0 +1,155 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as tp
+
+import exca
+import pydantic
+import pytorch_lightning as pl
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+import torch
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+
+
+class Model(pl.LightningModule):
+    def __init__(self, pretrained: bool, learning_rate: float = 0.001):
+        super(Model, self).__init__()
+        self.pretrained = pretrained
+        self.learning_rate = learning_rate
+        self.model = resnet18(pretrained=pretrained)
+        self.model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
+        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 10)
+        self.loss_fn = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x):
+        return self.model(x)
+
+    def _step(self, batch):
+        x, y = batch
+        y_hat = self(x)
+        loss = self.loss_fn(y_hat, y)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        loss = self._step(batch)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss = self._step(batch)
+        self.log("val_loss", loss)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
+
+
+class Data(pl.LightningDataModule):
+    def __init__(self, batch_size: int):
+        super().__init__()
+        self.batch_size = batch_size
+
+    def _dataloader(self, train: bool):
+        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
+        dset = datasets.MNIST('', train=train, download=True, transform=transform)
+        return torch.utils.data.DataLoader(dset, batch_size=self.batch_size)
+
+    def train_dataloader(self):
+        return self._dataloader(train=True)
+
+    def val_dataloader(self):
+        return self._dataloader(train=False)
+
+
+
+class ModelConfig(pydantic.BaseModel):
+    pretrained: bool = True
+    learning_rate: float = 0.001
+
+    model_config = pydantic.ConfigDict(extra="forbid")
+
+    def build(self) -> Model:
+        return Model(**self.dict())
+
+
+class DataConfig(pydantic.BaseModel):
+    batch_size: int = 64
+
+    model_config = pydantic.ConfigDict(extra="forbid")
+
+    def build(self) -> Data:
+        return Data(**self.dict())
+
+
+class TrainerConfig(pydantic.BaseModel):
+    max_epochs: tp.Optional[int] = None
+
+    model_config = pydantic.ConfigDict(extra="forbid")
+
+    def build(self, checkpoint_path: str | None = None) -> Trainer:
+        if checkpoint_path:
+            callbacks = [ModelCheckpoint(
+                dirpath=checkpoint_path,
+                save_top_k=1,
+                monitor="val_loss",
+                mode="min")
+            ]
+        else:
+            callbacks = None
+        return Trainer(**self.dict(), callbacks=callbacks)
+
+class Experiment(pydantic.BaseModel):
+    model: ModelConfig = ModelConfig()
+    data: DataConfig = DataConfig()
+    trainer: TrainerConfig = TrainerConfig()
+    infra: exca.TaskInfra = exca.TaskInfra(folder='.cache/')
+
+    @property
+    def checkpoint_path(self):
+        # Define the checkpoint directory
+        checkpoint_dir = self.infra.uid_folder() / 'checkpoint'
+
+        # Find the latest checkpoint if it exists
+        checkpoints = sorted(checkpoint_dir.glob('*.ckpt'))
+        ckpt_path = sorted(checkpoints)[-1] if checkpoints else None
+        return ckpt_path
+
+    @infra.apply
+    def fit(self):
+        # Configure
+        data = self.data.build()
+        model = self.model.build()
+        trainer = self.trainer.build(self.infra.folder)
+
+        # Fit model
+        trainer.fit(model, data, ckpt_path=self.checkpoint_path)
+
+        # Return model if not saved 
+        if self.checkpoint_path is None:
+            return model
+
+    def validate(self):
+        data = self.data.build()
+        model = self.model.build()
+        trainer = self.trainer.build(self.infra.folder)
+
+        trained_model = self.fit()
+        if trained_model is None:
+            trained_model = model.__class__.load_from_checkpoint(self.checkpoint_path)
+
+        return trainer.validate(trained_model, dataloaders=data.val_dataloader())
+
+
+if __name__ == '__main__':
+    config = dict(
+        model={'learning_rate': .01},
+        trainer={'max_epochs': 2},
+        infra={'folder': '.cache/'}
+    )
+    exp = Experiment(**config)
+    score = exp.validate()
+    print(score)

exca/steps/__init__.py

@@ -28,3 +28,5 @@
 from . import backends
 from .base import Chain as Chain
 from .base import Step as Step
+from .base import to_chain as to_chain
+from .base import to_step as to_step

exca/steps/base.py

@@ -28,6 +28,9 @@
 from . import backends
 from .backends import NoValue
 
+X = tp.TypeVar("X")
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -425,3 +428,210 @@ def clear_cache(self, recursive: bool = True) -> None:
                 step.clear_cache()
         if self.infra:
             self.infra.clear_cache()
+
+
+def _resolve_input_params(
+    func: tp.Callable[..., tp.Any],
+    input_params: tp.Sequence[str] | None,
+) -> tuple[str, ...]:
+    """Determine which function parameters are pipeline inputs.
+
+    If *input_params* is ``None``, parameters without a default value are
+    used.  Returns a validated tuple of parameter names.
+    """
+    sig = inspect.signature(func)
+    if input_params is None:
+        return tuple(
+            p.name for p in sig.parameters.values() if p.default is inspect._empty
+        )
+    for name in input_params:
+        if name not in sig.parameters:
+            raise ValueError(
+                f"input_params contains {name!r} which is not a parameter "
+                f"of {func.__name__}; available: {list(sig.parameters)}"
+            )
+    return tuple(input_params)
+
+
+def to_step(
+    func: tp.Callable[..., X],
+    input_params: tp.Sequence[str] | None = None,
+) -> tp.Type[Step]:
+    """Create a Step subclass from a function.
+
+    Parameters
+    ----------
+    func:
+        The function to wrap.  Parameters that are **not** pipeline inputs
+        must have type annotations and become pydantic model fields.
+    input_params:
+        Names of parameters that receive pipeline input at runtime.
+        By default (``None``), every parameter **without a default value**
+        is treated as a pipeline input.  Pass an explicit list (possibly
+        empty) to override.
+
+        - 0 input params  -> generator step, ``_run(self)``
+        - 1 input param   -> ``_run(self, value)``
+        - N input params  -> ``_run(self, value)`` where *value* is
+          unpacked as a tuple of length N
+
+    Example
+    -------
+    >>> def multiply(value: float, coeff: float = 2.0) -> float:
+    ...     return value * coeff
+    >>> MultiplyStep = to_step(multiply)   # value has no default -> input
+    >>> MultiplyStep(coeff=3).run(5.0)
+    15.0
+    """
+    resolved = _resolve_input_params(func, input_params)
+    resolved_set = set(resolved)
+    reserved = {"infra"}
+    sig = inspect.signature(func)
+
+    fields: dict[str, tp.Any] = {}
+    for p in sig.parameters.values():
+        if p.name in resolved_set:
+            continue
+        if p.name in reserved:
+            raise ValueError(
+                f"Parameter {p.name!r} is reserved and cannot be used " "as a field name"
+            )
+        if p.annotation in (tp.Any, inspect._empty):
+            raise ValueError(
+                f"Parameter {p.name!r} of {func.__name__} needs a type " "annotation"
+            )
+        default = ... if p.default is inspect._empty else p.default
+        fields[p.name] = (p.annotation, default)
+
+    # Build _run that calls the original function
+    if len(resolved) == 0:
+
+        def _run(self: tp.Any) -> tp.Any:  # type: ignore[misc]
+            return func(**{n: getattr(self, n) for n in fields})
+
+    elif len(resolved) == 1:
+        _input_name = resolved[0]
+
+        def _run(self: tp.Any, value: tp.Any) -> tp.Any:  # type: ignore[misc]
+            kwargs = {n: getattr(self, n) for n in fields}
+            kwargs[_input_name] = value
+            return func(**kwargs)
+
+    else:
+        _input_names = resolved
+
+        def _run(self: tp.Any, value: tp.Any) -> tp.Any:  # type: ignore[misc]
+            kwargs = {n: getattr(self, n) for n in fields}
+            kwargs.update(zip(_input_names, value))
+            return func(**kwargs)
+
+    Model: tp.Type[Step] = pydantic.create_model(
+        func.__name__ + "_Step",
+        **fields,
+        __base__=Step,
+        __module__=func.__module__,
+    )
+    Model._run = _run  # type: ignore[assignment]
+    return Model
+
+
+_FuncOrNamed = tp.Union[
+    tp.Callable[..., tp.Any],
+    tp.Tuple[str, tp.Callable[..., tp.Any]],
+]
+
+
+def to_chain(
+    *funcs: _FuncOrNamed,
+    infra: tp.Any = None,
+) -> tp.Type[Chain]:
+    """Create a Chain subclass from plain functions.
+
+    Each function is converted to a Step via :func:`to_step` and exposed
+    as a field on the returned Chain subclass, keyed by the function
+    name.  This lets you parameterize each step independently.
+
+    To use the same function more than once, pass ``(name, func)``
+    tuples to give each occurrence a distinct field name.
+
+    Parameters
+    ----------
+    *funcs:
+        Functions (or ``(name, func)`` tuples) to chain sequentially.
+    infra:
+        Optional default infra for the Chain (e.g.
+        ``{"backend": "Cached", "folder": "/tmp/cache"}``).
+
+    Example
+    -------
+    >>> def generate(seed: int = 42) -> float:
+    ...     import random; return random.Random(seed).random()
+    >>> def scale(x: float, factor: float = 10.0) -> float:
+    ...     return x * factor
+    >>> MyChain = to_chain(generate, ("upscale", scale), ("downscale", scale))
+    >>> chain = MyChain(upscale=dict(factor=100), downscale=dict(factor=0.5))
+    >>> result = chain.run()
+    """
+    if not funcs:
+        raise ValueError("to_chain requires at least one function")
+
+    # Normalise inputs to (name, func) pairs
+    named: list[tuple[str, tp.Callable[..., tp.Any]]] = []
+    for entry in funcs:
+        if isinstance(entry, tuple):
+            name, func = entry
+            named.append((name, func))
+        else:
+            named.append((entry.__name__, entry))
+
+    # Check for duplicate field names
+    seen: dict[str, int] = {}
+    for i, (name, _) in enumerate(named):
+        if name in seen:
+            first = named[seen[name]][1].__name__
+            current = named[i][1].__name__
+            raise ValueError(
+                f"Duplicate field name {name!r} (from {first!r} and "
+                f"{current!r}); use (name, func) tuples to disambiguate"
+            )
+        seen[name] = i
+
+    step_classes = [to_step(f) for _, f in named]
+    _field_names = tuple(name for name, _ in named)
+
+    # One field per step, typed as its Step subclass (with a default
+    # instance so all fields are optional).
+    model_fields: dict[str, tp.Any] = {}
+    for (name, _), StepCls in zip(named, step_classes):
+        model_fields[name] = (StepCls, StepCls())
+
+    # Override ``steps`` with an empty default; it is rebuilt from the
+    # per-function fields in model_post_init.
+    model_fields["steps"] = (tp.Sequence[Step], ())
+
+    _default_infra = infra
+
+    chain_name = "_".join(_field_names) + "_Chain"
+
+    Model: tp.Type[Chain] = pydantic.create_model(
+        chain_name,
+        **model_fields,
+        __base__=Chain,
+        __module__=named[0][1].__module__,
+    )
+
+    _super_post_init = Chain.model_post_init
+
+    def _model_post_init(self: tp.Any, __context: tp.Any) -> None:
+        if not self.steps:
+            built = tuple(getattr(self, name) for name in _field_names)
+            object.__setattr__(self, "steps", built)
+        if _default_infra is not None and self.infra is None:
+            infra_obj = backends.Backend.model_validate(_default_infra)
+            infra_obj._step = self
+            object.__setattr__(self, "infra", infra_obj)
+        _super_post_init(self, __context)
+
+    Model.model_post_init = _model_post_init  # type: ignore[assignment]
+
+    return Model