jejjohnson / jejjohnson/pipekit

pipekit-train[lightning]: implement Lightning backend adapter (v0.2)

Open
#13 0 comments 0 reactions 0 assignees View on GitHub
type:feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.