---
name: zyplib-experiment
description: Use when an agent needs to build, refactor, or review an experiment on top of `zyplib.experiment`. Trigger when the user asks to "use zyplib" to organize a training script, create an experiment runner, convert ad-hoc train code into `prepare_run` / `build_data` / `build_model` / `train_model`, define `TrialPlan`, or configure `ExperimentEffects` / `RecordConfig` / report export behavior.
---

# zyplib-experiment

Build the experiment from the agent's perspective: the user gives a training goal, and you must shape it into the current `zyplib.experiment` contract.

The source code of zyplib.experiment can be viewed in this way:
- First locate the installed/importable `zyplib` package root.
  - Example:

    ```bash
    python -c "import zyplib; from pathlib import Path; print(Path(zyplib.__file__).resolve().parent)"
    ```
- Then use Read/View tools on `<zyplib>/experiment/`.
  - Focus on `_types.py` for public types and contracts.
  - Focus on `context.py` for `ExperimentContext`, `TrialContext`, and record helpers.
  - Focus on `procedure.py` for the main execution flow, `run_experiment`, effects, seed schedule, and report export.
  - Focus on `runner.py` for the `BasicExperiment` class wrapper.

## Understand what the framework is for

`zyplib.experiment` is a run/trial orchestration layer.

Use it when one user request like "use zyplib to train/evaluate/repeat this model" needs to become:

- one run directory
- one or more trials inside that run
- stable per-trial execution steps
- aggregated final report output

Think in this order:

1. One **run** = one top-level experiment execution.
2. One run contains one or more **trials**.
3. Each trial runs the same lifecycle: `build_data -> build_model -> train_model`.
4. `prepare_run` prepares run-shared state before the trial loop starts.
5. `build_trial_plan` defines how many trials exist and what identifies each trial.

If the user only wants one training/evaluation pass, `build_trial_plan` is still valid; it usually just returns `TrialPlan.single(...)`.

## Start from the full skeleton

Map the target script into this structure before writing code:

```python
def prepare_run(ctx: ExperimentContext[Params]) -> Any:
    ...


def build_trial_plan(ctx: ExperimentContext[Params]) -> TrialPlan:
    ...


def build_data(ctx: TrialContext[Params]) -> ExperimentData:
    ...


def build_model(ctx: TrialContext[Params]) -> ModelSetup:
    ...


def train_model(
    ctx: TrialContext[Params],
    model_setup: ModelSetup,
    data: ExperimentData,
) -> TrialResult | dict[str, float]:
    ...


report = run_experiment(
    name='...',
    params=params,
    root_dir='...',
    prepare_run=prepare_run,
    plan=build_trial_plan,
    build_data=build_data,
    build_model=build_model,
    train_model=train_model,
    effects=ExperimentEffects(...),
)
```

Do not invent extra public lifecycle steps unless the codebase already requires them.

Simple mental model:

```text
prepare_run(ctx)      -> prepare shared runtime state
build_trial_plan(ctx) -> define the trial loop
for each trial:
  build_data(ctx)
  build_model(ctx)
  train_model(ctx, model_setup, data)
```

## Build each step deliberately

### 1. `prepare_run(ctx)`

Write run-level code only.

- Input: `ExperimentContext[Params]`
- Output: shared payload written to `ctx.shared`
- Write here:
  - run-level seed policy
  - precomputed split/fold schedule
  - reusable file index / cache / metadata
  - run-level notes via `ctx.set(...)`

Do not build dataloaders or models here.

### 2. `build_trial_plan(ctx)`

Define the trial schedule.

- Input: `ExperimentContext[Params]`
- Output: `TrialPlan`
- Write here:
  - `TrialPlan.single(...)`
  - `TrialPlan.repeat(...)`
  - `TrialPlan.kfold(...)` only when the script truly has fold semantics
  - `trial_seeds=...` when the real per-trial schedule is already known

Treat `trial_seed` as schedule metadata. The framework does not call `seed_everything()` for you.

### 3. `build_data(ctx)`

Construct the current trial's data.

- Input: `TrialContext[Params]`
- Output: `ExperimentData`
- Write here:
  - split selection based on `ctx.trial_seed`, `ctx.trial_index`, or `ctx.shared`
  - dataset / dataloader construction
  - serializable data metadata via `info={...}` (auto_record will consume this)
  - runtime-only extras via `extras={...}` (framework will not touch these)

`info` is for metadata that should appear in `report.yaml`. `extras` is for runtime objects the framework should ignore (e.g. extra domain loaders, preprocessing pipelines, dataset references).

This step must reflect actual trial differences. If all trials load the same split, the experiment is probably wrong.

### 4. `build_model(ctx)`

Construct model-bound training objects.

- Input: `TrialContext[Params]`
- Output: `ModelSetup`
- Write here:
  - model
  - optimizer
  - lr scheduler
  - model hparams

`ModelSetup` does not carry extras. Objects like `loss_fn` or custom hooks belong in `train_model`, where the user has full control.

Prefer returning `ModelSetup` instead of spreading model setup across multiple public hooks.

### 5. `train_model(ctx, model_setup, data)`

Run the backend-specific training/testing flow.

- Inputs: `TrialContext`, `ModelSetup`, `ExperimentData`
- Output: `TrialResult` or `dict[str, float]`
- Write here:
  - Fabric / torch / trainer calls
  - checkpoint loading / test pass
  - final metric collection

Put formal outputs in `TrialResult`. Do not rebuild a handwritten mega-report dict here.

## Put values in the correct container

- `TrialResult.metrics` for official numeric outputs
- `TrialResult.summary` for structured result summary
- `TrialResult.artifacts` for artifact paths
- `ExperimentData.info` for serializable data metadata (consumed by auto_record)
- `ExperimentData.extras` for runtime-only data objects (ignored by framework)
- `ctx.shared` for runtime-only shared objects across trials
- `ctx.set(key, value, namespace=...)` for small recorded metadata

If a value is large, non-serializable, or only useful during execution, prefer `ctx.shared` or `ExperimentData.extras`.

## Configure side effects explicitly

Use:

- `ExperimentEffects.capture_stdout`
- `ExperimentEffects.record_env`
- `RecordConfig`

Use `RecordConfig` to decide:

- which namespaces are injected at runtime
- what export rule controls report visibility
- whether framework auto-record is enabled

`export_rule` is the only export decision point. A list means namespace allowlist; a callable means full predicate logic.

## Read references only when needed

- Read `references/record-and-effects.md` when deciding how `ctx.set(...)`, `export_rule`, `snapshot_states()`, and auto-record should behave.
- Read `references/examples.md` when you need concrete construction patterns or want to sanity-check a script-to-experiment mapping.

## Finish with these checks

1. Does each step have the correct input/output shape?
2. Does `build_data()` actually consume the current trial identity?
3. Does `ctx.trial_seed` match the real per-trial schedule value when seeds matter?
4. Are effects and export behavior explicit?
5. Is `TrialResult` used for formal outputs instead of ad-hoc record blobs?
6. Is `ExperimentData.info` limited to serializable metadata? Are runtime objects in `extras`?
7. Does `build_model()` return only model-bound objects without smuggling extras?
