diff --git a/.github/workflows/test-type-lint.yaml b/.github/workflows/test-type-lint.yaml index bd267c29..b3f41601 100644 --- a/.github/workflows/test-type-lint.yaml +++ b/.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: | diff --git a/docs/infra/example_lightning.py b/docs/infra/example_lightning.py new file mode 100644 index 00000000..08426c2a --- /dev/null +++ b/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) diff --git a/exca/steps/__init__.py b/exca/steps/__init__.py index 747c7934..d027171d 100644 --- a/exca/steps/__init__.py +++ b/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 diff --git a/exca/steps/base.py b/exca/steps/base.py index 00b304a0..215622dd 100644 --- a/exca/steps/base.py +++ b/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 diff --git a/exca/steps/test_steps.py b/exca/steps/test_steps.py index c319b801..35aa953d 100644 --- a/exca/steps/test_steps.py +++ b/exca/steps/test_steps.py @@ -7,6 +7,7 @@ """Tests for Step and Chain basic functionality (no caching tests here, see test_cache.py).""" import pickle +import random import typing as tp from pathlib import Path @@ -16,7 +17,7 @@ import exca from . import backends, conftest -from .base import Chain, Input, Step +from .base import Chain, Input, Step, to_chain, to_step # ============================================================================= # Basic execution (no infra) @@ -312,3 +313,90 @@ def _forward(self, x: float) -> float: # .forward() call triggers its own warning with pytest.warns(DeprecationWarning, match="forward.*deprecated.*run"): assert step.forward(5.0) == 10.0 + + +# ============================================================================= +# to_step / to_chain helpers +# ============================================================================= + + +def generate(seed: int = 42) -> float: + return random.Random(seed).random() + + +def scale(x: float, factor: float = 10.0) -> float: + return x * factor + + +def test_to_step() -> None: + # generator (all defaults) + G = to_step(generate) + gen = G(seed=123) # type: ignore[call-arg] + assert gen._is_generator() + assert gen.run() == gen.run() + # transformer (auto-detect required param as input) + M = to_step(scale) + assert M(factor=3.0).run(5.0) == 15.0 # type: ignore[call-arg] + assert M().run(5.0) == 50.0 + + # explicit input_params override + def add(a: float, b: float) -> float: + return a + b + + AddStep = to_step(add, input_params=["a"]) + assert AddStep(b=10.0).run(5.0) == 15.0 # type: ignore[call-arg] + + # multiple inputs -> tuple unpacking + def combine(x: int, y: int, s: float = 1.0) -> float: + return (x + y) * s + + assert to_step(combine)(s=2.0).run((3, 7)) == 20.0 # type: ignore[call-arg] + # in a Chain + chain = Chain(steps=[G(seed=42), M(factor=100.0)]) # type: ignore[call-arg] + assert chain.run() == pytest.approx(random.Random(42).random() * 100.0) + + # validation errors + def f(x: int) -> int: + return x + + with pytest.raises(ValueError, match="not a parameter"): + to_step(f, input_params=["nope"]) + + def g(x: int, y) -> int: # type: ignore[no-untyped-def] + return x + y + + with pytest.raises(ValueError, match="needs a type annotation"): + to_step(g, input_params=["x"]) + + def h(infra: int) -> int: + return infra + + with pytest.raises(ValueError, match="reserved"): + to_step(h, input_params=[]) + + +def test_to_chain() -> None: + MyChain = to_chain(generate, scale) + assert issubclass(MyChain, Chain) + # defaults + assert MyChain().run() == pytest.approx( # type: ignore[call-arg] + random.Random(42).random() * 10.0 + ) + # custom params via dict + c = MyChain( # type: ignore[call-arg] + generate=dict(seed=123), scale=dict(factor=100.0) + ) + assert c.run() == pytest.approx(random.Random(123).random() * 100.0) + # partial override + c2 = MyChain(scale=dict(factor=5.0)) # type: ignore[call-arg] + assert c2.run() == pytest.approx(random.Random(42).random() * 5.0) + # (name, func) tuples for duplicate functions + MyChain2 = to_chain(generate, ("up", scale), ("down", scale)) + c3 = MyChain2(up=dict(factor=100.0), down=dict(factor=0.5)) # type: ignore[call-arg] + assert c3.run() == pytest.approx(generate() * 100.0 * 0.5) + # duplicate bare names rejected + with pytest.raises(ValueError, match="Duplicate"): + to_chain(scale, scale) + # empty rejected + with pytest.raises(ValueError, match="at least one"): + to_chain()