Add Dynamic Filter

Add filtering hooks in rollout and buffer stages while preserving sample-group contracts.

When to Use

Use this skill when:

• User asks to filter sample groups during dynamic sampling
• User asks to customize buffer extraction strategy
• User asks to mask/remove some rollout samples before training
• User asks to process all generated samples for logging/analysis

Step-by-Step Guide

Step 1: Pick the Correct Hook

• Dynamic sampling filter: --dynamic-sampling-filter-path
• Buffer filter: --buffer-filter-path
• Per-sample rollout filter: --rollout-sample-filter-path
• All-samples post process: --rollout-all-samples-process-path

Step 2: Implement the Function Signature

Dynamic sampling filter (called in vime/rollout/vllm_rollout.py):

def filter_function(args, samples, **kwargs):
    # return DynamicFilterOutput or bool

Preferred return type:

from vime.rollout.filter_hub.base_types import DynamicFilterOutput

return DynamicFilterOutput(keep=True, reason=None)

Buffer filter (called in vime/rollout/data_source.py):

def buffer_filter(args, rollout_id, buffer, num_samples):
    return selected_groups

Rollout sample filter:

def rollout_sample_filter(args, samples):
    # modify sample.remove_sample in-place where needed

All-samples process:

def process_all_samples(args, all_samples, data_source):
    ...

Step 3: Preserve Group Structure

• Keep list[list[Sample]] structure intact where required.
• Do not flatten sample groups unless downstream path expects flattened samples.
• For sample removal, prefer sample.remove_sample=True over deleting objects.

Step 4: Wire and Validate

Example wiring:

--dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
--buffer-filter-path <module>.buffer_filter
--rollout-sample-filter-path <module>.rollout_sample_filter
--rollout-all-samples-process-path <module>.process_all_samples

Step 5: Measure Side Effects

• Check final sample count remains aligned with rollout_batch_size expectations.
• Verify drop reasons are surfaced in rollout metrics when dynamic filter is used.
• Validate training still receives valid loss masks/rewards after filtering.

Common Mistakes

• Returning wrong container type for buffer filter
• Dropping samples by deletion instead of mask flagging
• Losing sample-group alignment in group-RM setup
• Adding expensive logic in hot filtering paths

Reference Locations

• Dynamic filter types: vime/rollout/filter_hub/base_types.py
• Dynamic filter example: vime/rollout/filter_hub/dynamic_sampling_filters.py
• Rollout generation hook points: vime/rollout/vllm_rollout.py
• Buffer filter hook point: vime/rollout/data_source.py

Add Eval Dataset Config

Configure evaluation datasets in vime with explicit dataset-level overrides and predictable runtime behavior.

When to Use

Use this skill when:

• User asks to add evaluation datasets for periodic eval
• User asks to migrate from --eval-prompt-data to structured --eval-config
• User asks for per-dataset eval overrides (sampling params, keys, rm_type, metadata)

Step-by-Step Guide

Step 1: Choose Config Entry Method

Supported inputs:

• Structured config file: --eval-config <yaml>
• Legacy CLI pairs: --eval-prompt-data <name1> <path1> <name2> <path2> ...

If --eval-interval is set, eval datasets must be configured.

Step 2: Build YAML with Required Fields

Each dataset needs at least:

• name
• path

Example:

eval:
  defaults:
    n_samples_per_eval_prompt: 1
    temperature: 0.7
    top_p: 1.0
  datasets:
    - name: aime
      path: /path/to/aime.jsonl
      rm_type: math
      input_key: prompt
      label_key: answer
      metadata_overrides:
        split: test

Step 3: Understand Override Priority

vime/utils/eval_config.py resolves fields in this order:

1. Dataset-level values in eval.datasets[*]
2. eval.defaults
3. CLI args fallback (for example eval_* or rollout_* fields)

Common overridable fields include:

• Runtime: n_samples_per_eval_prompt, temperature, top_p, top_k, max_response_len
• Sample keys: input_key, label_key, tool_key, metadata_key
• Extra: rm_type, custom_generate_function_path, metadata_overrides

Step 4: Wire Eval Function if Needed

By default, eval uses --eval-function-path (defaults to rollout function path). Use a separate eval function when inference/eval behavior must differ from training rollout.

Step 5: Validate Parsing and Runtime

• Start with config parsing sanity by running a short launch command.
• Confirm dataset entries are loaded into args.eval_datasets.
• Verify output keys match eval logging/metrics expectations.

Common Mistakes

• Missing name in dataset entries
• Odd-length --eval-prompt-data pairs
• Setting --eval-interval without any eval dataset
• Mixing reward dict outputs without eval_reward_key configuration

Reference Locations

• Eval config model: vime/utils/eval_config.py
• Eval config resolution: vime/utils/arguments.py
• Eval rollout path: vime/rollout/vllm_rollout.py
• Customization docs: docs/en/get_started/customization.md

Add Reward Function

Implement custom reward logic and connect it to vime rollout/training safely.

When to Use

Use this skill when:

• User asks to add new reward computation logic
• User asks to integrate an external reward service
• User asks to customize reward normalization/post-processing

Step-by-Step Guide

Step 1: Choose Reward Mode

Pick one of these:

• Single-sample mode (--group-rm disabled): custom function gets one Sample
• Group/batch mode (--group-rm enabled): custom function gets list[Sample]

vime.rollout.rm_hub.__init__.py calls your function via --custom-rm-path.

Step 2: Create Reward Module

Create vime/rollout/rm_hub/<your_rm>.py.

Supported signatures:

async def custom_rm(args, sample):
    return float_reward_or_reward_dict

async def custom_rm(args, samples):
    return list_of_rewards

If using group mode, return one reward per sample in input order.

Step 3: Keep Reward Type Consistent

• Return scalar numeric rewards unless your pipeline explicitly uses keyed rewards.
• If using reward dicts, ensure downstream reward_key / eval_reward_key is configured.
• Keep exceptions explicit for invalid metadata instead of silently returning zeros.

Step 4: Optional Reward Post-Processing

To customize normalization/shaping before advantage computation, add:

def post_process_rewards(args, samples):
    # return (raw_rewards, processed_rewards)
    ...

Wire with:

--custom-reward-post-process-path <module>.post_process_rewards

This hook is consumed in vime/ray/rollout.py.

Step 5: Wire and Validate

Use:

--custom-rm-path vime.rollout.rm_hub.<your_rm>.custom_rm

Common Mistakes

• Returning wrong output shape in group mode
• Mixing scalar rewards and reward dicts without reward_key config
• Doing blocking network calls without async handling
• Forgetting to validate reward behavior on truncated/failed samples

Reference Locations

• Reward dispatch: vime/rollout/rm_hub/__init__.py
• Reward post-process hook: vime/ray/rollout.py
• Customization docs: docs/en/get_started/customization.md

Add Rollout Function

Implement a custom rollout function and integrate it safely with vime training/eval flow.

When to Use

Use this skill when:

• User asks to add a new rollout task or rollout generation function
• User asks to replace default vime.rollout.vllm_rollout.generate_rollout
• User asks to customize train/eval data generation behavior

Step-by-Step Guide

Step 1: Choose the Right Starting Point

Start from one of these references:

• Async RL-style rollout: vime/rollout/vllm_rollout.py
• Simple SFT-style rollout: vime/rollout/sft_rollout.py

If the task needs engine-based async generation and rewards, use the vLLM path as base. If the task is file/buffer-driven and simple, use sft path as base.

Step 2: Create the New Rollout Module

Create a new file, for example: vime/rollout/<your_rollout>.py

Required callable signature:

def generate_rollout(args, rollout_id, data_source, evaluation=False) -> RolloutFnTrainOutput | RolloutFnEvalOutput:
    ...

Return types are defined in vime/rollout/base_types.py.

Step 3: Implement Train and Eval Branches Explicitly

• For training (evaluation=False), return RolloutFnTrainOutput(samples=..., metrics=...)
• For evaluation (evaluation=True), return RolloutFnEvalOutput(data=..., metrics=...)

Minimal skeleton:

from vime.rollout.base_types import RolloutFnTrainOutput, RolloutFnEvalOutput


def generate_rollout(args, rollout_id, data_source, evaluation=False):
    if evaluation:
        result = {
            "custom_eval": {
                "rewards": [],
                "truncated": [],
                "samples": [],
            }
        }
        return RolloutFnEvalOutput(data=result)

    groups = data_source.get_samples(args.rollout_batch_size)
    # fill Sample fields needed by training: tokens/response_length/reward/status (+ loss_mask when needed)
    return RolloutFnTrainOutput(samples=groups)

Step 4: Keep Data Contract Compatible

For each generated sample, ensure required training fields are populated consistently with your objective:

• tokens
• response_length
• reward (or reward dict if your setup uses keyed rewards)
• status

If partial rollout or masking logic is involved, keep loss_mask semantics consistent with existing behavior.

Step 5: Wire Through Arguments

Set your function path via CLI:

--rollout-function-path vime.rollout.<your_rollout>.generate_rollout

The default and signature expectation are documented in:

• vime/utils/arguments.py
• docs/en/get_started/customization.md

Common Mistakes

• Returning raw Python lists/dicts with mismatched schema in custom path
• Implementing only training branch and forgetting evaluation branch
• Generating samples without required fields (tokens, response_length, reward, status)
• Using blocking-heavy logic in high-frequency rollout paths without batching/concurrency control

Reference Locations

• Default rollout: vime/rollout/vllm_rollout.py
• Simple custom example: vime/rollout/sft_rollout.py
• Output dataclasses: vime/rollout/base_types.py
• Wiring/loading: vime/ray/rollout.py
• Argument definition: vime/utils/arguments.py
• Customization docs: docs/en/get_started/customization.md

Add Tests and CI

Add reliable tests and integrate them with vime CI flow.

When to Use

Use this skill when:

• User asks to add tests for new behavior
• User asks to fix or update existing tests in tests/
• User asks to update CI workflow behavior
• User asks how to run targeted checks before PR

Step-by-Step Guide

Step 1: Pick the Right Test Pattern

• Follow existing naming: tests/test_<feature>.py
• Start from nearest existing test file for your model/path
• Keep test scope small and behavior-focused

Step 2: Keep CI Compatibility

• CI executes registered test files with python tests/<file>.py, not only pytest discovery. New CPU pytest files should include:

import pytest

NUM_GPUS = 0

if __name__ == "__main__":
    raise SystemExit(pytest.main([__file__]))

• run-ci-changed extracts a top-level NUM_GPUS = <N> constant from added/modified tests/test_*.py and tests/plugin_contracts/test_*.py; if missing, it defaults to 8 GPUs. Set NUM_GPUS = 0 for CPU-only tests.
• For GPU/e2e tests, follow the nearby file pattern (prepare(), execute(), NUM_GPUS, and any model/dataset constants).

Step 3: Register Tests in GitHub CI

Whenever adding, moving, or renaming a test file, update the GitHub workflow template before finishing:

1. Add the test to the appropriate matrix in .github/workflows/pr-test.yml.j2.
   • CPU-only pytest/unit tests usually belong in cpu-unittest with num_gpus: 0.
   • GPU/e2e tests should be placed beside the nearest similar model/path test with the matching num_gpus and environment fields.
2. Regenerate workflows:

python .github/workflows/generate_github_workflows.py

3. Include both .github/workflows/pr-test.yml.j2 and the generated .github/workflows/pr-test.yml in the change set.

Only skip fixed matrix registration when the test is intentionally helper-only or manually invoked; state that reason in the final response.

Step 4: Run Local Validation

• Run the exact existing test files you changed, if any.
• For new registered tests, run the same shape CI will use, for example python tests/test_new_file.py.
• Run repository-wide checks only when they are already part of the task or workflow.
• Avoid documenting placeholder test commands that may not exist in the current tree.

Step 5: Keep Workflow Template as Source of Truth

For CI workflow changes unrelated to a new, moved, or renamed test:

1. Edit .github/workflows/pr-test.yml.j2
2. Regenerate workflows:

python .github/workflows/generate_github_workflows.py

3. Include both the template and generated workflow file in the change set (.j2 and .yml). If the user asked for a commit, commit both.

Step 6: Provide Verifiable PR Notes

Include:

• Which tests were added/changed
• Where each new/renamed test was registered in .github/workflows/pr-test.yml.j2
• Exact commands executed
• GPU assumptions for each test path
• Why this coverage protects against regression

Common Mistakes

• Editing generated workflow file only
• Relying on run-ci-changed discovery for a new test that should run in the regular PR matrix
• Forgetting NUM_GPUS = 0 on a CPU-only changed test, causing run-ci-changed to default to 8 GPUs
• Adding a CPU pytest file that passes under pytest tests/foo.py but fails under CI's python tests/foo.py
• Adding tests without following existing constants/conventions
• Making tests too large or non-deterministic
• Skipping local validation and relying only on remote CI

Reference Locations

• Pytest config: pyproject.toml
• Tests: tests/
• CI template: .github/workflows/pr-test.yml.j2
• CI guide: docs/en/developer_guide/ci.md

Vime Code Review Preferences

Apply these lightweight review heuristics when changing vime code.

Prefer Direct APIs Over Thin Wrappers

• Remove helper layers that only rename a call, format one path, or forward arguments without owning meaningful behavior.
• Prefer calling the concrete reusable API directly, for example a *_to_path helper when the caller already knows the destination path.
• Keep a wrapper only if it owns a real boundary: compatibility, validation, nontrivial error policy, lifecycle management, metrics/logging semantics, async/retry behavior, or cross-module ownership.
• Avoid moving a redundant wrapper's body into another file just to preserve the wrapper shape. Inline the simple call at the natural ownership site.
• When removing a wrapper, search for sibling wrappers and nearby helpers with rg and delete confirmed dead functions in the same pass.
• Treat single-use convenience functions as suspicious when their only job is path formatting plus forwarding. Prefer the caller owning that one line.

Make Branches Explain Themselves

• Order conditionals by semantic precedence: special transport/lifecycle modes first, then explicit mode choices, then default paths.
• Prefer predicates that fully describe the branch, such as mode == "full" and transport == "disk", over a broad predicate followed by an assert that explains what the branch really meant.
• Use asserts as invariants for impossible states after validation, not as a substitute for clear branch conditions.

Keep Abstractions Honest

• Add an abstraction only when it removes real duplication, hides fragile mechanics, or clarifies ownership.
• When a review comment points out repeated indirection, look for a smaller public surface rather than adding another alias.
• Preserve existing behavior intentionally. If cleanup changes error handling, logging, or failure visibility, call that out in the final response.