---
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 should reflect the real source of trial differences.

- If trials differ by data split / fold, consume `ctx.trial_seed`, `ctx.trial_index`, or `ctx.shared` here.
- If trials intentionally reuse the same split and only vary model init / training randomness / hyperparameters, that is also valid; make that choice explicit in code or recorded metadata.

### 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

If the trial only needs to return flat numeric metrics, returning `dict[str, float]` is valid. Use `TrialResult` when you also need structured `summary` or `artifacts`. 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 export decision point for recorded namespace state (`TrialReport.records`). It does not control `TrialResult.metrics`, `TrialResult.summary`, `TrialResult.artifacts`, or the final aggregated report. 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. Are formal outputs returned either as `dict[str, float]` for simple metrics or `TrialResult` for richer 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?
