jejjohnson / jejjohnson/pipekit
pipekit-train[lightning]: implement Lightning backend adapter (v0.2)
- Dominant language
- Python
- Stars
- 0
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Problem / Request
Implement the Lightning backend adapter in `pipekit_train.adapters.lightning` so PyTorch users can drive `TrainingLoop` end-to-end. The module currently raises `NotImplementedError` from `run(loop)`.
## User Story
> As a PyTorch user, I want `TrainingLoop(backend="lightning", ...)` to actually train my LightningModule via Lightning's `Trainer.fit`, so I can use the same pipekit-train surface I'd otherwise use with the Equinox backend.
## Motivation
- Broadens v0.1's user base from JAX-only to the (much larger) PyTorch Lightning ecosystem.
- The full design lives in `packages/pipekit-train/docs/design/api/adapters.md` (Lightning section) — the planned implementation is "synthesize a `LightningModule` + `LightningDataModule`, then `Trainer.fit`".
- Adapter pattern is well-trodden (Equinox adapter is the reference impl); same shape, different backend.
## Proposed API
```python
# pipekit_train/adapters/lightning.py
def run(loop: TrainingLoop) -> tuple[Operator, dict[str, Any]]:
"""Train ``loop`` with PyTorch Lightning.
Returns:
(trained_model_op, backend_info) where trained_model_op is a
pipekit-array.ModelOp (or inline _LightningModelOp wrapper)
carrying the trained nn.Module.
"""
```
## Design Snapshot
From `docs/design/api/adapters.md` (Lightning section):
> 1. Synthesise a `LightningModule` from `loop.model_op` + `loop.loss` (or accept a user-supplied `task` that is itself a `LightningModule` factory).
> 2. Synthesise a `LightningDataModule` from `loop.dataset` / `loop.val_dataset`.
> 3. Build a `Trainer` with `max_steps=loop.max_steps`, mapping pipekit callbacks to Lightning's `pl.Callback` subclass surface (e.g. `LogToExperiment → MLFlowLogger` or `WandbLogger`).
> 4. Call `trainer.fit(module, datamodule)`.
> 5. Wrap the final module's `nn.Module` in `pipekit-array.ModelOp` and return.
`pipekit-array.ModelOp` doesn't exist yet — for v0.2 the adapter can ship an inline `_LightningModelOp(Operator)` wrapper (mirroring the v0.1 `EquinoxModelOp` pattern), then swap when `pipekit-array` lands.
## References & Existing Code
- Reference adapter: `packages/pipekit-train/src/pipekit_train/adapters/equinox.py`
- Adapter contract: `docs/design/architecture.md` §5 (Backend adapter contract)
- Existing stub: `packages/pipekit-train/src/pipekit_train/adapters/lightning.py`
- Existing scaffold tests: `packages/pipekit-train/tests/test_adapters.py` (these need to flip from "expect NotImplementedError" to actual coverage)
## Implementation Steps
- [ ] `_build_lightning_module(loss, model_op) → LightningModule`
- [ ] `_build_datamodule(dataset, val_dataset, batch_size) → LightningDataModule`
- [ ] `_callback_bridge` — translate pipekit `Callback` hooks to `pl.Callback`
- [ ] `_optimizer_from_config(optimizer_config) → torch.optim.Optimizer`
- [ ] `_LightningModelOp(Operator)` — inline wrapper for the trained `nn.Module`
- [ ] `run(loop)` — orchestration (mirrors `adapters/equinox.run`)
- [ ] Tests under `tests/adapters/test_lightning.py` gated by `pytest.importorskip("lightning")`
- [ ] End-to-end test: train tiny PyTorch MLP on synthetic regression; assert loss decreases ≥50%
- [ ] Flip `tests/test_adapters.py::test_lightning_*` from "expects NotImplementedError" to "expects success when [lightning] extra installed; expects ImportError otherwise"
## Definition of Done
- [ ] `TrainingLoop(backend="lightning").run()` trains an MLP and returns `(trained_op, artifact)`.
- [ ] All four pre-commit gates pass with `[lightning]` extra installed.
- [ ] API docs page (`docs/api/pipekit-train.md`) adds Lightning section.
## Relationships
- Parent (theme epic): pipekit-train v0.2
- Related: `pipekit-array` for the future `ModelOp` (not blocking)
Contributor guide
Assessment
This issue has not been assessed yet.