diff --git a/README.md b/README.md index 11d1de4..8de5405 100644 --- a/README.md +++ b/README.md @@ -176,8 +176,7 @@ with a generated dataset card. ### Reusing a fixed population `--locale` samples personas fresh each run. To fix the locales and per-locale -row counts up front (for a matched-pair comparison, say), build a panel and -pass it instead: +row counts up front, build a panel and pass it instead: ```bash usersim panel --locale en_US --num-personas 50 --out panel.parquet @@ -187,6 +186,19 @@ usersim simulate --panel panel.parquet --out output/trajectories Note that personas are re-sampled per run either way; a panel sets the shape of the run, not its cast. +To fix the cast as well, store the resolved rows once and run those: + +```bash +usersim simulate --locale en_US --num-rows 50 --materialize-inputs --out rows.jsonl +usersim simulate --inputs rows.jsonl --models assistant_a.toml --out output/trajectories +usersim simulate --inputs rows.jsonl --models assistant_b.toml --out output/trajectories +``` + +Both runs meet the same simulated users in the same scenarios. Each setup gets +its own run under `--out`, and each row its own `trajectory_id`; the stored +row's id is kept as `input_id`, so the two runs pair on it. Running again with +the same setup resumes its run. + ### Models Inference defaults to [build.nvidia.com](https://build.nvidia.com) through diff --git a/docs/engine/EXTERNAL_PROBE_RUNTIME.md b/docs/engine/EXTERNAL_PROBE_RUNTIME.md index 7cfe954..2a67479 100644 --- a/docs/engine/EXTERNAL_PROBE_RUNTIME.md +++ b/docs/engine/EXTERNAL_PROBE_RUNTIME.md @@ -47,6 +47,73 @@ usersim simulate --locale en_US --num-rows 1 \ --materialize-inputs --out episode-inputs.jsonl ``` +## Running stored rows + +To run stored rows with your configured models, without hosting the episode, +pass them back to `simulate`. Each row runs exactly as stored, and is saved +where `usersim eval` reads it: + +```bash +usersim simulate --inputs episode-inputs.jsonl --models my_models.toml --out output/trajectories +``` + +or from Python: + +```python +from pathlib import Path + +from usersim.engine.external import simulate_episode_inputs + +result = simulate_episode_inputs(rows, out=Path("output/trajectories"), models_path=Path("my_models.toml")) +print(result.run_id, result.written, result.skipped) +``` + +- **One run per setup.** A setup is everything that can change a + conversation: + - each simulator model as Data Designer sends it, provider settings + included + - the rows' shared settings + - the asset contents + - the `USERSIM_*` variables + - probe prompt versions + - extension, User Sim and Data Designer releases + + Running again with the same setup resumes its run and skips the rows already + saved there, so a caller can run rows in batches. A different setup starts + its own run. +- **Ids.** Each saved row gets a `trajectory_id` of its own, derived from the + stored row and the setup. The stored row's id is kept as `input_id`, so two + runs of the same rows pair on it: `x.merge(y, on="input_id")`. A stored row + whose contents changed under the same id is refused in the run that already + holds it. +- **Stored fields are saved with the rows.** A field whose values are all one + plain type (text, true/false, a decimal, or an integer within signed 64 + bits) keeps that type; anything else, such as a nested value, integers mixed + with decimals, or a larger integer, is saved as JSON text. The choice is + made from the rows of each call, so running the same rows again saves each + field the same way. Batches into one run read back exactly when every batch + makes the same choice for a field. When two don't (integers in one, + decimals in another), reading the run reconciles them, which can round + large integers, turn true/false into numbers, or mix plain text with JSON + text. Run such a field's rows in one call. +- **The conversation runs with the stored id,** as a hosted episode does, so + every setup sees the same scenario (the tools `tool_calling` offers, say). A + move hook registered during such a run therefore reports the stored id, the + saved row's `input_id`. +- **Settings come from the rows.** Each row runs with the settings it was + stored with, so a sampling flag (`--max-turns`, `--probe-mix`, ...) set to + anything but its default is refused. The rows in one call must share their + settings apart from locale and asset folder. + `--assets-dir` replaces the asset folder the rows were stored with, and is + required when that folder doesn't exist on this machine. +- **Concurrent writers.** Writes into `--out` take a lock, so two calls can + share it. On a filesystem without file locks, or whose locks don't reach + every node (Lustre mounted without `flock` or with `localflock`, NFS with + `local_lock=flock`, `local_lock=all` or `nolock`), writing goes ahead, with + a warning where file locks are missing. There, two calls at the same time + with one setup can each start a run, and two resumes of one run at the same + time can save a row twice, or both versions of a changed row. + ## Declaring the host's models Trajectory identity is keyed on the resolved model id, and `identity_disclosure` diff --git a/docs/engine/README.md b/docs/engine/README.md index b7a228a..66a37a9 100644 --- a/docs/engine/README.md +++ b/docs/engine/README.md @@ -237,7 +237,7 @@ usersim simulate --locale en_US --num-rows 5 \ Each probe resolves its curated asset at simulator-call time under `assets//`: `assets/sov_ai_facts//sample.yaml`, `assets/sov_ai_dynamic//categories.yaml`, `assets/sov_ai_multilingual_parity/sample.yaml`, `assets/safety_chat_pressure/sample.yaml`, `assets/safety_agentic/{sample.yaml,action_taxonomy.yaml}`, and `assets/identity_disclosure/spec.yaml`. Override paths via `USERSIM_SOV_AI_FACTS_BANK_` / `USERSIM_SOV_AI_DYNAMIC_TAXONOMY_` / `USERSIM_SOV_AI_MULTILINGUAL_PARITY_BANK` / `USERSIM_SAFETY_CHAT_PRESSURE_BANK` / `USERSIM_SAFETY_AGENTIC_BANK` / `USERSIM_SAFETY_AGENTIC_ACTION_TAXONOMY` / `USERSIM_IDENTITY_DISCLOSURE_SPEC` env vars when running against private banks; an identity spec named by its override can extend the shipped one rather than replace it. -**Panel-driven matrix sampling**: the `sov_ai_multilingual_parity`, `safety_chat_pressure`, `safety_agentic` and `identity_disclosure` probes all accept optional input-row columns that override their persona-derived defaults: `query_id` for `sov_ai_multilingual_parity`, `pressure_strategy_id` + `target_request_id` for `safety_chat_pressure`, `action_request_id` for `safety_agentic`, `identity_tactic_id` + `identity_strategy_id` for `identity_disclosure`. To use them, build the panel with the override column populated (typically via `usersim panel` followed by a small pandas script that adds the column) and pass `--panel ` to `usersim simulate` instead of `--locale`. Without overrides, every probe falls back to its persona-derived default: useful for smoke runs. +**Pinning a probe's scenario with override columns**: the `safety_chat_pressure`, `safety_agentic` and `identity_disclosure` probes read optional columns from their input row that override their persona-derived choice: `pressure_strategy_id` + `target_request_id` for `safety_chat_pressure`, `action_request_id` for `safety_agentic`, `identity_tactic_id` + `identity_strategy_id` for `identity_disclosure`. A panel cannot carry them, since `usersim simulate --panel` reads only a panel's locales and row counts. Instead, write the rows with `usersim simulate --materialize-inputs`, set the column on the rows to pin, and run them with `usersim simulate --inputs` (see [Running stored rows](EXTERNAL_PROBE_RUNTIME.md#running-stored-rows)). Without an override, each probe falls back to its persona-derived choice. `sov_ai_multilingual_parity` takes no override: its query follows from the persona and the seed. ## Trajectory Evaluator diff --git a/docs/outputs.md b/docs/outputs.md index 9c96688..56805ad 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -13,6 +13,7 @@ simulated conversation, each carrying: - `simulation_outcome` (JSON) holds a row-level structured record: status, failure-class taxonomy, per-actor attribution, per-model token counts, wall-clock, provenance (code SHA + asset bank versions + prompt versions). - `simulation_traces` (JSON): per-turn / per-call telemetry (judge-gate ratings, fourth-wall pre-filter triggers, capitulation checks, tool-call verification, etc.). - A content-hashed `persona_uuid` (stable across runs) and a deterministic `trajectory_id` (hash of `persona + probe + resolved model identities + prompt version`). These are the simulator's idempotency keys. +- In a run of stored rows (`usersim simulate --inputs`), `trajectory_id` is instead a hash of the stored row and of the setup it ran with: models as sent, the rows' shared settings, asset contents, `USERSIM_*` variables, and prompt and package versions (see [Running stored rows](engine/EXTERNAL_PROBE_RUNTIME.md#running-stored-rows)). `input_id` holds the stored row's own `trajectory_id`, so two setups run over the same stored rows pair on it. A row keeps the `usersim_provenance` it was materialized with, while the run's manifest records the code and versions the run itself used. > **Reasoning-trace export:** `usersim select` and its Hugging Face upload path > retain `conversation_messages`, so captured `reasoning_content` travels with diff --git a/src/usersim/cli/_inputs.py b/src/usersim/cli/_inputs.py new file mode 100644 index 0000000..ed15cf0 --- /dev/null +++ b/src/usersim/cli/_inputs.py @@ -0,0 +1,664 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Running stored episode rows: ``usersim simulate --inputs``. + +Rows written by ``usersim simulate --materialize-inputs`` run exactly as they +were stored, so two setups can be compared on the same simulated users. + +- **Identity.** Each run of a row gets its own ``trajectory_id``: a hash of the + stored row and of everything about the setup that can change the + conversation (see :func:`setup_fingerprint`). The stored row's own id is + saved as ``input_id``, the column that pairs two runs of the same rows. +- **The conversation.** It runs with the stored id as ``trajectory_id``, as + the hosted runtime does, so whatever a probe derives from its id (the tools + ``tool_calling`` offers, say) is the same for every setup. The new id is + written only when the rows are saved. +- **Runs.** A run holds one setup. A resume reopens the latest run of stored + rows whose setup fingerprint matches; otherwise the first saved rows claim + a new run. Every save happens under the root's write lock, after checking + the run's saved rows again. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import sys +import tempfile +import time +from collections import Counter +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +logger = logging.getLogger("usersim.cli.simulate") + +#: The saved row's copy of the stored row's ``trajectory_id``. +INPUT_ID_COLUMN = "input_id" + +#: Names the id formula, so a later formula never collides with this one. +_ID_SCHEME = "usersim-inputs-v1" + +#: ``USERSIM_*`` variables that cannot change a conversation: debug logging, +#: where the models file is (the models themselves are fingerprinted), +#: provenance labels, and the evaluator's audit settings. Every other +#: ``USERSIM_*`` variable is part of the setup, including ones added later. +_OPERATIONAL_ENV = frozenset( + { + "USERSIM_DEBUG_LOG", + "USERSIM_MODELS_CONFIG", + "USERSIM_CODE_SHA", + "USERSIM_NEMOTRON_PERSONAS_VERSION", + "USERSIM_AUDIT_MODEL", + "USERSIM_AUDIT_REASONING_EFFORT", + } +) + +#: Settings that may differ between the rows of one invocation. +_PER_ROW_SETTINGS = ("locale", "assets_dir") + + +@dataclass(frozen=True) +class EpisodeInputs: + """Stored rows checked and ready to run.""" + + rows: list[dict[str, Any]] + #: The ``usersim_config`` every row shares, apart from ``_PER_ROW_SETTINGS``. + config: dict[str, Any] + #: The asset root the rows run against. + assets_dir: Path + #: The file the rows came from, when they came from one. + source: Path | None = None + + +@dataclass(frozen=True) +class InputsRun: + """What one call to :func:`simulate_inputs_to_dataset` did.""" + + run_id: str | None + resumed: bool + written: int + skipped: int + + +# ── loading and checking ──────────────────────────────────────────────── + + +def load_episode_inputs(path: str | Path) -> list[dict[str, Any]]: + """Read stored rows from a ``.jsonl`` or ``.parquet`` file, values as stored. + + Parquet is read through pyarrow, not pandas, so no value is re-typed on + the way in. + """ + path = Path(path) + if path.suffix == ".jsonl": + with path.open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + if path.suffix == ".parquet": + import pyarrow.parquet as pq + + return pq.read_table(path).to_pylist() + raise ValueError(f"expected a .jsonl or .parquet file, got {path}") + + +def _plain(value: Any) -> Any: + """Plain Python for a row built in memory, which may hold numpy values.""" + if isinstance(value, Mapping): + return {key: _plain(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + if isinstance(value, str): + return str(value) + if hasattr(value, "tolist"): + return _plain(value.tolist()) + return value + + +def prepare_episode_inputs( + rows: list[Mapping[str, Any]], + *, + assets_dir: str | Path | None = None, + source: str | Path | None = None, +) -> EpisodeInputs: + """Check stored rows before any model call, and settle what they run against. + + Raises ``ValueError`` naming the row and the problem. The rows must share + one setup: their ``usersim_config`` may differ only in locale and asset + root, so that one run holds one setup. + """ + from pydantic import ValidationError + + from usersim.engine.config import ConversationSimulatorConfig + from usersim.engine.core.episode_input import EPISODE_INPUT_COLUMN + from usersim.engine.core.probes import resolve_probe + + if not rows: + raise ValueError("no rows to run") + checked: list[dict[str, Any]] = [] + seen: set[str] = set() + for i, raw in enumerate(rows): + if not isinstance(raw, Mapping): + raise ValueError(f"row {i} is not a mapping") + row = _plain(dict(raw)) + stored_id = row.get("trajectory_id") + if not isinstance(stored_id, str) or not stored_id: + raise ValueError(f"row {i} has no trajectory_id; each stored row needs its own") + where = f"row {i} ({stored_id})" + if stored_id in seen: + raise ValueError(f"{where}: trajectory_id appears in more than one row; each stored row needs its own") + seen.add(stored_id) + for reserved in (INPUT_ID_COLUMN, EPISODE_INPUT_COLUMN): + if reserved in row: + raise ValueError(f"{where} already has {reserved!r}; that is a run's output, not a stored row") + if not isinstance(row.get("usersim_config"), Mapping): + raise ValueError(f"{where} has no usersim_config; run rows written by --materialize-inputs") + try: + config = ConversationSimulatorConfig.model_validate(dict(row["usersim_config"])) + except ValidationError as error: + raise ValueError(f"{where}: usersim_config does not load: {error}") from None + if not row.get(config.persona_column): + raise ValueError(f"{where} has no {config.persona_column}") + try: + resolve_probe(str(row.get(config.probe_type_column))) + except KeyError as error: + raise ValueError(f"{where}: {error.args[0]}") from None + if "locale" in row and row["locale"] != config.locale: + raise ValueError(f"{where}: locale {row['locale']!r} differs from its config's {config.locale!r}") + try: + json.dumps(row, allow_nan=True) + except (TypeError, ValueError) as error: + raise ValueError(f"{where} holds a value JSON cannot represent: {error}") from None + checked.append(row) + return EpisodeInputs( + rows=checked, + config=_shared_config(checked), + assets_dir=_resolve_assets_dir(checked, assets_dir), + source=Path(source).resolve() if source is not None else None, + ) + + +def _shared_config(rows: list[dict[str, Any]]) -> dict[str, Any]: + """The setup every row shares; ``ValueError`` naming the settings that differ.""" + + def setup(row: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in row["usersim_config"].items() if key not in _PER_ROW_SETTINGS} + + first = setup(rows[0]) + for row in rows[1:]: + other = setup(row) + differing = sorted(key for key in first.keys() | other.keys() if first.get(key) != other.get(key)) + if differing: + shown = ", ".join(f"{key}: {first.get(key)!r} vs {other.get(key)!r}" for key in differing) + raise ValueError( + f"the rows were stored with different settings ({shown}); one run holds one setup, " + "so run them separately" + ) + return dict(rows[0]["usersim_config"]) + + +def _resolve_assets_dir(rows: list[dict[str, Any]], given: str | Path | None) -> Path: + """``--assets-dir`` when given, else the one asset root the rows were stored with.""" + from usersim.engine.core._assets import packaged_assets_dir + + if given is not None: + path = Path(given).resolve() + if not path.is_dir(): + raise ValueError(f"--assets-dir {path} is not a directory") + return path + stored = {row["usersim_config"].get("assets_dir") for row in rows} + hint = f"pass --assets-dir (the assets that ship with User Sim are at {packaged_assets_dir()})" + if len(stored) != 1 or None in stored: + raise ValueError(f"the rows do not name one asset folder; {hint}") + path = Path(stored.pop()) + if not path.is_dir(): + raise ValueError(f"the asset folder the rows were stored with, {path}, does not exist here; {hint}") + return path + + +# ── identity ──────────────────────────────────────────────────────────── + + +def _digest(value: Any) -> str: + """SHA-256 of ``value``'s canonical JSON form.""" + text = json.dumps(value, sort_keys=True, ensure_ascii=False, separators=(",", ":"), default=str) + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _contents_digest(path: Path) -> str: + """SHA-256 over a file's bytes, or over every file under a folder by relative path. + + Hidden files and folders (``.DS_Store``, ``.ipynb_checkpoints``) are left + out: no loader reads them, and they change without the assets changing. + """ + if path.is_file(): + return hashlib.sha256(path.read_bytes()).hexdigest() + entries = sorted( + item + for item in path.rglob("*") + if item.is_file() and not any(part.startswith(".") for part in item.relative_to(path).parts) + ) + return _digest( + [[item.relative_to(path).as_posix(), hashlib.sha256(item.read_bytes()).hexdigest()] for item in entries] + ) + + +def _model_settings(models: Any, aliases: list[str]) -> tuple[dict[str, Any], dict[str, Any]]: + """Each alias's settings as Data Designer sends them: readable, and as hashed. + + The readable form names a provider's extra headers without their values, + which may be credentials; the hashed form carries a digest of them. + """ + from usersim.cli._models import resolved_model_providers, to_model_configs + + configs = {config.alias: config for config in to_model_configs(models)} + providers = {provider.name: provider for provider in resolved_model_providers(models)} + readable: dict[str, Any] = {} + hashed: dict[str, Any] = {} + for alias in sorted(set(aliases)): + config = configs.get(alias) + if config is None: + readable[alias] = hashed[alias] = None + continue + params = config.inference_parameters + settings = {"model": config.model, "request": dict(params.generate_kwargs), "timeout": params.timeout} + provider = providers.get(config.provider) + if provider is None: + readable[alias] = hashed[alias] = {**settings, "provider": None} + continue + served = { + "endpoint": provider.endpoint, + "provider_type": provider.provider_type, + "extra_body": provider.extra_body or {}, + } + headers = provider.extra_headers or {} + readable[alias] = {**settings, "provider": {**served, "extra_headers": sorted(headers)}} + hashed[alias] = {**settings, "provider": {**served, "extra_headers": _digest(headers)}} + return readable, hashed + + +def _asset_roots_digest(assets_dir: Path) -> list[str]: + """The contents of every asset root a run searches, in priority order. + + Every root counts, not just the first copy of each file: an identity spec + that ``extends`` its own probe reads the copy below it on the path. Where a + root is does not count, so the same assets in another folder are the same + setup. + """ + from usersim.engine.core._assets import asset_search_path, reset_runtime_assets_dir, set_runtime_assets_dir + + token = set_runtime_assets_dir(assets_dir) + try: + roots = asset_search_path() + finally: + reset_runtime_assets_dir(token) + return [_contents_digest(root) if root.exists() else "missing" for root in roots] + + +def _environment_settings() -> dict[str, list[str]]: + """Every ``USERSIM_*`` variable that can change a conversation. + + A part of a value (split on ``os.pathsep``) that names an existing file or + folder counts by its contents, so the same bank in another place is the + same setup, and an edited bank is a different one. + """ + settings: dict[str, list[str]] = {} + for name in sorted(os.environ): + if not name.startswith("USERSIM_") or name in _OPERATIONAL_ENV: + continue + parts = os.environ[name].split(os.pathsep) + settings[name] = [_contents_digest(Path(part)) if part and Path(part).exists() else part for part in parts] + return settings + + +def _prompt_versions() -> dict[str, str]: + """``PROMPT_VERSION`` of every registered probe, read as the episode preamble reads it.""" + from usersim.engine.core.probes import known_probes, resolve_probe + + return { + label: str(getattr(sys.modules.get(resolve_probe(label).__module__), "PROMPT_VERSION", "v1.0")) + for label in known_probes() + } + + +def _extension_identities() -> list[dict[str, Any]]: + """Every entry point User Sim discovers, with the package release behind it.""" + from importlib.metadata import entry_points + + from usersim.engine.core._extensions import ENTRY_POINT_GROUPS, extensions_disabled + + if extensions_disabled(): + return [] + identities: list[dict[str, Any]] = [] + for group in ENTRY_POINT_GROUPS: + for entry in sorted(entry_points(group=group), key=lambda e: (e.name, e.value)): + dist = getattr(entry, "dist", None) + identities.append( + { + "group": group, + "name": entry.name, + "value": entry.value, + "distribution": dist.name if dist else None, + "version": dist.version if dist else None, + } + ) + return identities + + +def _data_designer_version() -> str | None: + """The installed Data Designer release, which shapes every model request.""" + from usersim.engine.core.manifest import _pkg_version + + return _pkg_version("data_designer") + + +def setup_fingerprint(models: Any, inputs: EpisodeInputs) -> tuple[str, dict[str, Any]]: + """The fingerprint of everything about a setup that can change a conversation. + + Returns the fingerprint and the run record that names it. Covered: + + - every simulator model alias as Data Designer sends it, provider settings + included (``api_key``, a provider's name and ``max_parallel_requests`` + are left out: they don't change what a model says); + - the shared ``usersim_config`` as loaded, defaults applied; + - the contents of every asset root, and every non-operational + ``USERSIM_*`` variable; + - every probe's ``PROMPT_VERSION``, every extension's package release, and + the User Sim and Data Designer releases. + + A code change that bumps none of these versions is not seen. + """ + from usersim.engine.config import MODEL_ALIASES, ConversationSimulatorConfig + from usersim.engine.core.manifest import _own_version + from usersim.engine.core.storage import RUN_KIND_STORED + + simulation = ConversationSimulatorConfig.model_validate(dict(inputs.config)).model_dump(mode="json") + for key in _PER_ROW_SETTINGS: + simulation.pop(key, None) + aliases = [*MODEL_ALIASES, simulation.get("finance_embedding_model_alias") or ""] + readable, hashed = _model_settings(models, [alias for alias in aliases if alias]) + fingerprint = _digest( + { + "models": hashed, + "simulation": simulation, + "assets": _asset_roots_digest(inputs.assets_dir), + "environment": _environment_settings(), + "prompts": _prompt_versions(), + "extensions": _extension_identities(), + "usersim": _own_version(), + "data_designer": _data_designer_version(), + } + ) + record = {"kind": RUN_KIND_STORED, "fingerprint": fingerprint, "models": readable, "simulation": simulation} + return fingerprint, record + + +def replay_trajectory_id(row: Mapping[str, Any], fingerprint: str) -> str: + """The ``trajectory_id`` of ``row`` run under the setup ``fingerprint`` names. + + The stored row is hashed whole, apart from its machine-specific asset + path, so a row whose contents changed under the same stored id is a + different conversation. + """ + stored = dict(row) + stored["usersim_config"] = {key: value for key, value in row["usersim_config"].items() if key != "assets_dir"} + composite = f"{_ID_SCHEME}|{_digest(stored)}|{fingerprint}" + return hashlib.sha256(composite.encode("utf-8")).hexdigest()[:16] + + +# ── running ───────────────────────────────────────────────────────────── + + +def _matching_run(out: Path, fingerprint: str) -> str | None: + """The latest run of stored rows made with this setup.""" + from usersim.engine.core.storage import RUN_KIND_STORED, latest_run_where + + def matches(run_id: str, record: dict[str, Any] | None) -> bool: + return bool(record) and record.get("kind") == RUN_KIND_STORED and record.get("fingerprint") == fingerprint + + return latest_run_where(out, matches) + + +def _saved_inputs(out: Path, run_id: str) -> dict[str, set[str]]: + """Each ``input_id`` saved in ``run_id``, with the trajectory ids saved under it.""" + from usersim.engine.core.storage import read_partitioned_dataset + + saved = read_partitioned_dataset(out, run=run_id, columns=["trajectory_id", INPUT_ID_COLUMN]) + pairs: dict[str, set[str]] = {} + if not saved.empty: + for input_id, trajectory_id in zip(saved[INPUT_ID_COLUMN].astype(str), saved["trajectory_id"].astype(str)): + pairs.setdefault(input_id, set()).add(trajectory_id) + return pairs + + +def _unsaved(planned: list[tuple[str, str]], saved: dict[str, set[str]], run_id: str) -> tuple[set[str], int]: + """The ``(input_id, trajectory_id)`` pairs not yet saved in ``run_id``, and how many were. + + Raises ``ValueError`` for an input saved there under a different id: its + stored contents changed, and one run must not hold two versions of it. + """ + todo: set[str] = set() + done = 0 + for input_id, trajectory_id in planned: + prior = saved.get(input_id) + if not prior: + todo.add(input_id) + elif trajectory_id in prior: + done += 1 + else: + raise ValueError( + f"stored row {input_id!r} already ran in run {run_id} with different contents; give the " + "changed row its own trajectory_id, or write to another --out" + ) + return todo, done + + +@dataclass(frozen=True) +class InputsPlan: + """What a run of stored rows would do, worked out before any model call.""" + + fingerprint: str + #: The record a new run with this setup is published with. + record: dict[str, Any] + #: The run of stored rows with this setup that a resume reopens, if any. + matching_run: str | None + #: Stored ids not yet saved in that run. + pending: set[str] + #: Rows already saved there. + saved: int + + +def plan_inputs_run(inputs: EpisodeInputs, models: Any, out: str | Path, *, skip_existing: bool = True) -> InputsPlan: + """Fingerprint the setup and find what the matching run already holds. + + Raises ``ValueError`` if that run holds one of the stored rows with + different contents. + """ + fingerprint, record = setup_fingerprint(models, inputs) + stored_ids = {row["trajectory_id"] for row in inputs.rows} + matching = _matching_run(Path(out).resolve(), fingerprint) if skip_existing else None + if matching is None: + return InputsPlan(fingerprint, record, matching_run=None, pending=stored_ids, saved=0) + planned = [(row["trajectory_id"], replay_trajectory_id(row, fingerprint)) for row in inputs.rows] + pending, saved = _unsaved(planned, _saved_inputs(Path(out).resolve(), matching), matching) + return InputsPlan(fingerprint, record, matching_run=matching, pending=pending, saved=saved) + + +def _column_config(inputs: EpisodeInputs, locale: str): + """The simulator column for ``locale``'s rows, on the run's asset root.""" + from usersim.engine.config import ConversationSimulatorConfig + + return ConversationSimulatorConfig.model_validate( + {**inputs.config, "locale": locale, "assets_dir": str(inputs.assets_dir)} + ) + + +#: Value types a stored column added back to the saved rows keeps as they are. +_PLAIN_TYPES = frozenset({str, bool, int, float}) + + +def _json_columns(rows: list[dict[str, Any]]) -> set[str]: + """The stored columns a run saves as JSON text, chosen from all of an invocation's rows. + + A column is saved as it is when every value present is of one plain type + (``str``, ``bool``, ``float``, or an ``int`` within 64 bits). Anything else, + a nested value or a mix such as an ``int`` beside a ``float``, would be + converted or refused by Arrow, so it is saved as JSON text, which keeps + every value exactly. The choice depends only on the stored rows, so every + invocation over the same rows, a resume included, saves a column the same + way. + """ + kinds: dict[str, set[type]] = {} + oversized: set[str] = set() + for row in rows: + for key, value in row.items(): + if value is None: + continue + kinds.setdefault(key, set()).add(type(value)) + if type(value) is int and not -(2**63) <= value < 2**63: + oversized.add(key) + return {key for key, types in kinds.items() if key in oversized or len(types) != 1 or not types <= _PLAIN_TYPES} + + +def _as_saved_rows(df, stored: dict[str, tuple[dict[str, Any], str]], json_columns: set[str]): + """The rows a run saves. + + The rows get the new ids, the stored id as ``input_id``, and a plain run's + columns. ``df`` comes back from Data Designer with each stored row's own + id as ``trajectory_id``. The carrier column and the full persona are + dropped, as a plain run drops the persona, and a stored column the run did + not keep is added back from the stored rows, as JSON text when it is one + of ``json_columns``. + """ + import pandas as pd + + from usersim.engine.core.episode_input import EPISODE_INPUT_COLUMN + + df = df.copy() + input_ids = df["trajectory_id"].astype(str).tolist() + df[INPUT_ID_COLUMN] = input_ids + df["trajectory_id"] = [stored[input_id][1] for input_id in input_ids] + df = df.drop(columns=[EPISODE_INPUT_COLUMN, "persona"], errors="ignore") + missing = sorted( + {key for row, _ in stored.values() for key in row} - set(df.columns) - {"usersim_config", "persona"} + ) + for key in missing: + values = [stored[input_id][0].get(key) for input_id in input_ids] + if key in json_columns: + values = [None if value is None else json.dumps(value, ensure_ascii=False) for value in values] + df[key] = pd.Series(values, index=df.index, dtype=object) + return df + + +def simulate_inputs_to_dataset( + *, + inputs: EpisodeInputs, + models: Any, + out: str | Path, + models_path: str | Path | None = None, + skip_existing: bool = True, + plan: InputsPlan | None = None, +) -> InputsRun: + """Run stored rows with ``models`` and save them as one run under ``out``. + + Each locale's rows run as one Data Designer batch, through the simulator + column ``usersim simulate`` uses. With ``skip_existing`` (the default), the + rows go into the latest run with this setup, skipping rows already saved + there; otherwise they start a new run. ``plan`` is + :func:`plan_inputs_run`'s result for the same arguments, when the caller + already has it. + + Raises ``ValueError`` when the run already holds a stored row with + different contents: before any model call when it is found up front, and + when the rows are saved when another writer saved them meanwhile, where + the write lock is shared. Where ``out`` is on a filesystem without file + locks, the rows are written without a lock shared between processes, + with a warning. Two calls at once with one setup can then each start a + run, and resuming one run from two places at once can save a row twice, + or both versions of a changed row. + """ + from usersim.cli._pipeline import build_inputs_config_builder + from usersim.engine.core.manifest import write_simulator_manifest + from usersim.engine.core.storage import ( + publish_new_run, + run_started_at, + run_write_lock, + write_locale_partition, + ) + + started_at = time.time() + out = Path(out).resolve() + out.mkdir(parents=True, exist_ok=True) + # A first look at the matching run, to skip saved rows before any model + # call and refuse changed ones early. Only the check under the lock counts. + plan = plan or plan_inputs_run(inputs, models, out, skip_existing=skip_existing) + fingerprint, matching, pending, skipped = plan.fingerprint, plan.matching_run, plan.pending, plan.saved + record = {**plan.record, "started_at": started_at} + stored = {row["trajectory_id"]: (row, replay_trajectory_id(row, fingerprint)) for row in inputs.rows} + logger.info("setup %s: %d stored rows, %d already saved", fingerprint[:12], len(stored), skipped) + + by_locale: dict[str, dict[str, tuple[dict[str, Any], str]]] = {} + for input_id in sorted(pending): + row, new_id = stored[input_id] + by_locale.setdefault(row["usersim_config"]["locale"], {})[input_id] = (row, new_id) + + # The rows go into the run the plan found, even if a newer run with this + # setup appears meanwhile, so they never split across two runs. + run_id: str | None = matching + resumed = matching is not None + json_columns = _json_columns(inputs.rows) + written = 0 + requested: Counter[str] = Counter() + for locale, group in by_locale.items(): + requested[locale] += len(group) + data_designer, builder = build_inputs_config_builder( + models=models, config=_column_config(inputs, locale), rows=[row for row, _ in group.values()] + ) + # Data Designer names a batch's artifacts by the second it started, so + # two invocations in one directory could share them; each batch gets + # its own folder, removed once its rows are loaded. + with tempfile.TemporaryDirectory(prefix="usersim_inputs_") as artifacts: + result = data_designer.create(builder, num_records=len(group), artifact_path=artifacts) + df = _as_saved_rows(result.load_dataset(), group, json_columns) + with run_write_lock(out): + target = run_id or (_matching_run(out, fingerprint) if skip_existing else None) + if target is not None: + planned = list(zip(df[INPUT_ID_COLUMN].astype(str), df["trajectory_id"].astype(str))) + todo, done = _unsaved(planned, _saved_inputs(out, target), target) + df = df[df[INPUT_ID_COLUMN].astype(str).isin(todo)] + skipped += done + # This invocation keeps the run it found, even if every row + # was saved there meanwhile. + resumed = resumed or run_id is None + run_id = target + if df.empty: + continue + if target is None: + run_id = publish_new_run(df, out, locale=locale, record=record) + else: + write_locale_partition(df, out, locale=locale, run_id=target) + written += len(df) + logger.info("locale=%s: wrote %d trajectories from stored rows into run=%s", locale, len(df), run_id) + + if written and run_id is not None: + probes = Counter(str(row.get("probe_type")) for row, _ in stored.values()) + try: + write_simulator_manifest( + run_id=run_id, + started_at_epoch=run_started_at(out, run_id) or int(started_at), + trajectory_root=out, + models_config=models, + models_config_path=Path(models_path) if models_path else None, + sim_config=_column_config(inputs, next(iter(requested))), + locales_requested=sorted(requested), + probe_mix={probe: count / len(stored) for probe, count in sorted(probes.items())}, + panel_path=inputs.source, + assets_dir=inputs.assets_dir, + trajectories_requested=dict(requested), + ) + except Exception as error: # pragma: no cover - never fail a finished run on manifest issues + logger.warning("failed to write run manifest: %s", error) + return InputsRun(run_id=run_id, resumed=resumed, written=written, skipped=skipped) diff --git a/src/usersim/cli/_models.py b/src/usersim/cli/_models.py index c55b55c..751c24f 100644 --- a/src/usersim/cli/_models.py +++ b/src/usersim/cli/_models.py @@ -642,6 +642,27 @@ def to_data_designer_kwargs(config: ModelsConfig) -> dict[str, Any]: return {"model_providers": merged} +def resolved_model_providers(config: ModelsConfig) -> list[Any]: + """The providers Data Designer serves this config's models through. + + The config's own providers merged with the defaults, or, when it + declares none, the defaults Data Designer loads for itself, from the + user's providers file. A provider's ``extra_body`` and ``extra_headers`` + are merged into every request it serves. + """ + providers = to_data_designer_kwargs(config).get("model_providers") + if providers is not None: + return providers + from data_designer.config.default_model_settings import get_builtin_model_providers, get_default_providers + + try: + return get_default_providers() + except FileNotFoundError: + # Data Designer writes its providers file, seeded with the built-ins, + # the first time it runs. + return get_builtin_model_providers() + + def to_model_configs(config: ModelsConfig) -> list[Any]: """Convert ``ModelsConfig`` into a list of ``dd.ModelConfig`` objects. diff --git a/src/usersim/cli/_pipeline.py b/src/usersim/cli/_pipeline.py index 31a5a8b..407cb20 100644 --- a/src/usersim/cli/_pipeline.py +++ b/src/usersim/cli/_pipeline.py @@ -566,6 +566,53 @@ def materialize_episode_inputs( return rows +def build_inputs_config_builder( + *, + models: ModelsConfig, + config: Any, + rows: list[dict[str, Any]], +): + """Build a Data Designer pipeline that runs stored episode rows as they are. + + Each row becomes one seed row: the columns the simulator column requires + (``persona``, ``probe_type``, ``theme``), plus the whole stored row as JSON + in ``EPISODE_INPUT_COLUMN``, which the generator unpacks before the probe + is built, so every value keeps the type it was stored with. ``config`` is + the simulator column, restored from the rows' ``usersim_config``. + + Returns ``(data_designer, config_builder)``, like + :func:`build_simulator_config_builder`. + """ + import data_designer.config as dd + import pandas as pd + from data_designer.interface import DataDesigner + + from usersim.engine.core.episode_input import EPISODE_INPUT_COLUMN + + def _text(value: Any) -> str: + return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + + seed = pd.DataFrame( + [ + { + config.persona_column: _text(row[config.persona_column]), + config.probe_type_column: str(row[config.probe_type_column]), + config.theme_column: _text(row.get(config.theme_column, "")), + EPISODE_INPUT_COLUMN: json.dumps( + {key: value for key, value in row.items() if key != "usersim_config"}, ensure_ascii=False + ), + } + for row in rows + ] + ) + data_designer = DataDesigner(**to_data_designer_kwargs(models)) + _set_run_config(data_designer, dd) + config_builder = dd.DataDesignerConfigBuilder(model_configs=to_model_configs(models)) + config_builder.with_seed_dataset(dd.DataFrameSeedSource(df=seed)) + config_builder.add_column(config) + return data_designer, config_builder + + def build_evaluator_config_builder( *, models: ModelsConfig, diff --git a/src/usersim/cli/panel.py b/src/usersim/cli/panel.py index c4dfcf0..22287f1 100644 --- a/src/usersim/cli/panel.py +++ b/src/usersim/cli/panel.py @@ -3,19 +3,13 @@ """``usersim panel`` — build a frozen persona seed parquet. -Why a panel? The simulator's identity scheme content-hashes the persona -(``persona_uuid``) so two runs that draw the same persona always -produce the same id. Persisting the panel separately from the simulator -run gives you: - -- **Replayability** — re-run the simulator against the same panel after - fixing a bug or swapping the assistant model; every ``trajectory_id`` - remains stable up to the bits that intentionally changed. -- **Matched-pair evaluation** — feed the same panel to two different - ``assistant_model`` configurations and join their trajectories on - ``persona_uuid`` for paired statistical tests. -- **Audit** — stash the panel parquet alongside a model release as - evidence of the population the model was evaluated against. +A panel fixes the shape of a run: ``usersim simulate --panel`` takes its +locales and per-locale row counts, and samples personas afresh. So a panel +does not replay a run, and two runs from one panel are not guaranteed to +meet the same people. To run the same simulated users under two setups, +write them once with ``usersim simulate --materialize-inputs`` and run them +with ``usersim simulate --inputs`` under each; the two runs pair on +``input_id``. The panel is a thin parquet with one row per persona: ``persona`` (JSON dict) plus the seven derived ``persona_*`` expression columns diff --git a/src/usersim/cli/simulate.py b/src/usersim/cli/simulate.py index dc520e7..7c9a7bf 100644 --- a/src/usersim/cli/simulate.py +++ b/src/usersim/cli/simulate.py @@ -21,6 +21,9 @@ Idempotent re-runs: if the output dataset already exists, rows whose ``trajectory_id`` is present are skipped (the simulator re-emits identical bytes for them anyway, but skipping avoids the LLM cost). + +``--inputs`` runs stored rows from ``--materialize-inputs`` exactly as +stored instead of sampling, one run per setup; see ``cli/_inputs.py``. """ from __future__ import annotations @@ -29,6 +32,7 @@ import logging import tempfile import time +from collections import Counter from fractions import Fraction from pathlib import Path @@ -64,6 +68,20 @@ def register(subparsers: argparse._SubParsersAction) -> None: action="append", help=("Locale to sample inline (repeatable). Use with --num-rows. Mutually exclusive with --panel."), ) + src.add_argument( + "--inputs", + type=Path, + help=( + "Run the stored rows in this .jsonl or .parquet file (from " + "--materialize-inputs) exactly as stored, instead of sampling, so " + "two setups meet the same simulated users. Each row runs with the " + "settings it was stored with, so a sampling flag (--probe-mix, " + "--max-turns, --random-seed, ...) set to anything but its default " + "is refused. Each setup gets its " + "own run under --out, and each row a trajectory_id of its own; " + "the stored id is kept as input_id, which pairs two runs." + ), + ) p.add_argument( "--num-rows", @@ -229,7 +247,27 @@ def register(subparsers: argparse._SubParsersAction) -> None: default="local", help="Where the run executes (default: local, in this process).", ) - p.set_defaults(func=run) + p.set_defaults( + func=run, + _stored_row_flag_defaults={dest: p.get_default(dest) for dest in _STORED_ROW_FLAGS}, + ) + + +#: Flags that set what ``--inputs`` takes from the stored rows instead, by +#: argparse destination. +_STORED_ROW_FLAGS = { + "num_rows": "--num-rows", + "max_turns": "--max-turns", + "max_assistant_attempts": "--max-assistant-attempts", + "probe_mix": "--probe-mix", + "random_seed": "--random-seed", + "finance_tier_mix": "--finance-tier-mix", + "finance_retrieval_mode": "--finance-retrieval-mode", + "no_store_reasoning": "--no-store-reasoning", + "toolset_seed_path": "--toolset-seed-path", + "match_persona_language": "--match-persona-language", + "materialize_inputs": "--materialize-inputs", +} def _parse_probe_mix(spec: str) -> dict: @@ -288,6 +326,9 @@ def run(args: argparse.Namespace) -> int: models = load_models_config(models_path) require_aliases(models, required=REQUIRED_SIMULATOR_ALIASES, context="`usersim simulate`") + if args.inputs is not None: + return _run_inputs(args, models, models_path) + probe_mix = _parse_probe_mix(args.probe_mix) # Resolve --assets-dir once, here, so every downstream consumer sees a @@ -435,6 +476,101 @@ def _materialize_inputs(args, models, probe_mix: dict[str, float]) -> int: return 0 +def _run_inputs(args: argparse.Namespace, models, models_path: Path) -> int: + """``usersim simulate --inputs``: run stored rows exactly as stored.""" + from usersim.cli._inputs import ( + load_episode_inputs, + plan_inputs_run, + prepare_episode_inputs, + simulate_inputs_to_dataset, + ) + from usersim.cli._models import warn_if_missing_api_key + + defaults = args._stored_row_flag_defaults + given = [flag for dest, flag in _STORED_ROW_FLAGS.items() if getattr(args, dest) != defaults[dest]] + if given: + raise SystemExit( + f"--inputs runs each row with the settings it was stored with; drop {', '.join(given)} " + "(to change them, materialize the rows again)" + ) + try: + inputs = prepare_episode_inputs( + load_episode_inputs(args.inputs), assets_dir=args.assets_dir, source=args.inputs + ) + plan = plan_inputs_run(inputs, models, args.out, skip_existing=not args.no_skip_existing) + except (OSError, ValueError) as error: + raise SystemExit(f"--inputs {args.inputs}: {error}") from None + + config = inputs.config + print("usersim simulate --inputs" + (" — dry run" if args.dry_run else "")) + print(f" inputs: {args.inputs} ({len(inputs.rows)} rows)") + print(f" locales: {sorted({row['usersim_config']['locale'] for row in inputs.rows})}") + print(f" probes: {dict(sorted(Counter(str(row['probe_type']) for row in inputs.rows).items()))}") + print( + " stored with: " + + ", ".join( + f"{key}={config.get(key)!r}" + for key in ("max_turns", "max_assistant_attempts", "random_seed", "store_reasoning") + ) + ) + print(f" assets_dir: {inputs.assets_dir}") + print(f" models config: {models_path}") + print(f" setup: {plan.fingerprint[:12]}") + if plan.matching_run is not None: + print(f" run_id: {plan.matching_run} (resume; {plan.saved} of {len(inputs.rows)} rows already saved)") + else: + print(" run_id: new, assigned when the first rows are saved") + print(f" output: {args.out}") + if args.dry_run: + return 0 + + if (msg := warn_if_missing_api_key(models)) is not None: + logger.warning(msg) + + def job() -> int: + try: + result = simulate_inputs_to_dataset( + inputs=inputs, + models=models, + out=args.out, + models_path=models_path, + skip_existing=not args.no_skip_existing, + plan=plan, + ) + except ValueError as error: + # Another writer saved one of these rows with different contents + # while this run was simulating. + raise SystemExit(f"--inputs {args.inputs}: {error}") from None + if not result.written: + print("no new trajectories generated") + return 0 + print(f"wrote {result.written} new trajectories to run={result.run_id} under {args.out}") + print( + f"\nNext step: score them\n usersim eval --run {result.run_id} \\\n" + f" --trajectories {args.out} --out output/evaluations" + ) + return 0 + + from usersim.engine.core.backends import JobSpec, get_backend + + backend = get_backend(getattr(args, "backend", "local")) + params = { + "inputs": str(args.inputs), + "out": str(args.out), + "models_path": str(models_path), + "assets_dir": str(inputs.assets_dir), + "skip_existing": not args.no_skip_existing, + } + handle = backend.submit(JobSpec(name="simulate-inputs", params=params, callable_=job)) + status = backend.poll(handle) + if status.error is not None: + raise status.error + if status.state == "failed" and status.exit_code is None: + logger.error("simulate --inputs failed on backend %r: %s", backend.name, status.message) + return 1 + return status.exit_code or 0 + + def _simulate_to_parquet( *, models, diff --git a/src/usersim/engine/core/episode_input.py b/src/usersim/engine/core/episode_input.py index 4590ccf..0998adb 100644 --- a/src/usersim/engine/core/episode_input.py +++ b/src/usersim/engine/core/episode_input.py @@ -31,6 +31,27 @@ if TYPE_CHECKING: from usersim.engine.config import ConversationSimulatorConfig +#: A seed column that carries a stored episode row whole, as JSON. A seed table +#: types each of its columns, so an ``int`` column with a ``float`` in another +#: row would reach the probe as ``float``; inside one JSON cell every value +#: keeps the type it was stored with. +EPISODE_INPUT_COLUMN = "usersim_episode_input" + + +def unpack_stored_episode(data: dict[str, Any]) -> None: + """Replace ``data``'s columns with the stored episode it carries, if any. + + The carrier column is removed, and the stored values overwrite the seed + table's copies of the same columns. + """ + cell = data.pop(EPISODE_INPUT_COLUMN, None) + if cell is None: + return + stored = json.loads(cell) if isinstance(cell, str) else cell + if not isinstance(stored, Mapping): + raise ValueError(f"{EPISODE_INPUT_COLUMN} must hold a JSON object, got {type(stored).__name__}") + data.update(stored) + @dataclass(frozen=True) class EpisodePreamble: diff --git a/src/usersim/engine/core/storage.py b/src/usersim/engine/core/storage.py index 536820c..1192561 100644 --- a/src/usersim/engine/core/storage.py +++ b/src/usersim/engine/core/storage.py @@ -219,6 +219,8 @@ def resolve_run_or_raise( #: The kind a record gives a run of rows the simulator sampled. A resume #: reopens only a run of its own kind. RUN_KIND_SAMPLED: str = "sampled" +#: The kind a record gives a run of stored rows that ``simulate --inputs`` ran. +RUN_KIND_STORED: str = "stored" #: A new run is assembled in a hidden folder with this prefix, then renamed. _CLAIM_PREFIX: str = ".claim-" @@ -314,7 +316,9 @@ def run_write_lock(root: str | Path) -> Iterator[Path]: Without a shared lock, a new run still never shares a folder with another, and every write has a part file of its own. What is lost: - resuming one run from two places at once can save a row twice. + resuming one run from two places at once can save a row twice, and two + writers that start a run of stored rows with one setup at once can make + two runs of it. """ from filelock import FileLock @@ -408,7 +412,8 @@ def _warn_without_flock(root: Path) -> None: _WARNED_WITHOUT_FLOCK.add(str(root)) logger.warning( "file locks (flock) are not supported under %s, so runs there are written without a lock " - "shared between processes. Resuming one run from two places at once can save a row twice.", + "shared between processes. Resuming one run from two places at once can save a row twice, " + "and two writers that start a run of stored rows with one setup at once can make two runs.", root, ) diff --git a/src/usersim/engine/external.py b/src/usersim/engine/external.py index c1e7d8f..ed82b1a 100644 --- a/src/usersim/engine/external.py +++ b/src/usersim/engine/external.py @@ -40,6 +40,7 @@ "construct_episode_preamble", "construct_probe_episode", "materialize_episode_inputs", + "simulate_episode_inputs", ] @@ -83,3 +84,50 @@ def materialize_episode_inputs( finance_tier_mix=finance_tier_mix, finance_retrieval_mode=finance_retrieval_mode, ) + + +def simulate_episode_inputs( + rows: list[dict[str, object]], + *, + out: Path, + models_path: Path | None = None, + assets_dir: Path | None = None, + skip_existing: bool = True, +): + """Run stored rows with the configured models and save them as a run under ``out``. + + The library form of ``usersim simulate --inputs``. ``rows`` are what + :func:`materialize_episode_inputs` returns, or the same rows read back + from storage. Each row runs exactly as stored, and is saved where + ``usersim eval`` reads it, with a ``trajectory_id`` of its own and the + stored id as ``input_id``. + + Each setup (models, settings, assets) gets its own run under ``out``. + Calling again with the same setup resumes that run and skips the rows it + already holds, so a caller can run a large set in batches (a stored field + that two batches save differently does not read back exactly; see + "Running stored rows" in ``docs/engine/EXTERNAL_PROBE_RUNTIME.md``). Raises + ``ValueError`` for rows that cannot run, before any model call, and for a + stored row that the run already holds with different contents: before any + model call when it is found up front, and when the rows are saved when + another writer saved it meanwhile, where the write lock is shared. Where + ``out`` is on a filesystem without file locks, the rows are written + without a lock shared between processes, with a warning. Two calls at + once with one setup can then each start a run, and resuming one run from + two places at once can save a row twice, or both versions of a changed + row. Returns an ``InputsRun`` with the run id and the number of rows + written and skipped. + """ + from usersim.cli._inputs import prepare_episode_inputs, simulate_inputs_to_dataset + from usersim.cli._models import REQUIRED_SIMULATOR_ALIASES, default_models_path, load_models_config, require_aliases + + path = models_path or default_models_path() + models = load_models_config(path) + require_aliases(models, required=REQUIRED_SIMULATOR_ALIASES, context="simulate_episode_inputs") + return simulate_inputs_to_dataset( + inputs=prepare_episode_inputs(rows, assets_dir=assets_dir), + models=models, + out=out, + models_path=path, + skip_existing=skip_existing, + ) diff --git a/src/usersim/engine/generator.py b/src/usersim/engine/generator.py index de920f5..a078c1e 100644 --- a/src/usersim/engine/generator.py +++ b/src/usersim/engine/generator.py @@ -26,7 +26,11 @@ ConversationSimulatorConfig, ) from usersim.engine.core._assets import reset_runtime_assets_dir, set_runtime_assets_dir -from usersim.engine.core.episode_input import EpisodeConstructionError, construct_probe_episode +from usersim.engine.core.episode_input import ( + EpisodeConstructionError, + construct_probe_episode, + unpack_stored_episode, +) from usersim.engine.core.llm import ( ContextWindowError, flush_debug_log, @@ -160,6 +164,9 @@ async def _generate_row(self, data: dict) -> dict: t_record_start = time.monotonic() + # A replayed row carries its stored episode whole; the probe is built + # from those values, not from the seed table's typed copies. + unpack_stored_episode(data) construction_error: EpisodeConstructionError | None = None try: constructed = construct_probe_episode(data, config=cfg, models=models) diff --git a/tests/engine/test_stored_episode_input.py b/tests/engine/test_stored_episode_input.py new file mode 100644 index 0000000..e1c59e6 --- /dev/null +++ b/tests/engine/test_stored_episode_input.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A row can carry its stored episode whole, as one JSON cell. + +The invariant: the probe is built from the stored values exactly, types +included, whatever a seed table would have made of them as separate columns. +""" + +from __future__ import annotations + +import json +from contextlib import ExitStack +from typing import Any +from unittest.mock import patch + +import pytest + +from usersim.engine import generator as generator_module +from usersim.engine.config import ConversationSimulatorConfig +from usersim.engine.core.episode_input import EPISODE_INPUT_COLUMN +from usersim.engine.generator import ConversationSimulatorGenerator + +_LLM_SITES = ( + "usersim.engine.core.llm", + "usersim.engine.core.simulation", + "usersim.engine.core.judges", + "usersim.engine.core.context", + "usersim.engine.core.translation", +) + +_STORED = { + "persona": {"first_name": "Sarah", "last_name": "Johnson", "age": 35, "city": "Seattle", "state": "WA"}, + "probe_type": "general_open_ended", + "theme": "cooking at home", + "trajectory_id": "stored-a1", + "ext_big": 9007199254740993, + "ext_ratio": 0.5, + "ext_flag": False, + "ext_nested": {"memory": [{"week": 1, "note": "slept badly"}], "score": 2}, +} + + +async def _stub_llm(models, alias, msgs, **kwargs): + if alias == "judge_model": + return {"role": "assistant", "content": "fine\nsuccess"} + if alias == "summary_model": + return {"role": "assistant", "content": "no"} + return {"role": "assistant", "content": "Thanks, that is helpful."} + + +class _Unused: + async def acompletion(self, *args, **kwargs): + raise AssertionError("acall_llm is patched") + + +async def _build_from(row: dict[str, Any]) -> dict[str, Any]: + """Run ``row`` and return the row the probe was built from.""" + seen: list[dict[str, Any]] = [] + construct = generator_module.construct_probe_episode + + def spy(data, **kwargs): + seen.append(dict(data)) + return construct(data, **kwargs) + + config = ConversationSimulatorConfig(name="conversation_messages", locale="en_US", max_turns=1, random_seed=1) + aliases = ("user_model", "assistant_model", "judge_model", "summary_model", "api_response_model") + generator = ConversationSimulatorGenerator.with_models(config, {alias: _Unused() for alias in aliases}) + with ExitStack() as stack: + stack.enter_context(patch.object(generator_module, "construct_probe_episode", spy)) + for site in _LLM_SITES: + stack.enter_context(patch(f"{site}.acall_llm", side_effect=_stub_llm)) + await generator.agenerate(row) + (built,) = seen + return built + + +def _seed_row(stored: dict[str, Any]) -> dict[str, Any]: + """The row as a seed table hands it over: the required columns, plus the stored episode.""" + return { + "persona": json.dumps(stored["persona"]), + "probe_type": stored["probe_type"], + "theme": stored["theme"], + EPISODE_INPUT_COLUMN: json.dumps(stored), + } + + +async def test_the_probe_is_built_from_the_stored_values_exactly(): + built = await _build_from(_seed_row(_STORED)) + + for key, value in _STORED.items(): + assert built[key] == value and type(built[key]) is type(value), key + + +async def test_the_stored_episode_wins_over_the_seed_columns(): + seed = {**_seed_row(_STORED), "theme": "a theme the seed table rewrote"} + + built = await _build_from(seed) + + assert built["theme"] == "cooking at home" + + +async def test_the_carrier_column_does_not_reach_the_probe(): + built = await _build_from(_seed_row(_STORED)) + + assert EPISODE_INPUT_COLUMN not in built + + +@pytest.mark.parametrize("cell", ["[1, 2]", '"text"']) +async def test_a_carrier_that_is_not_a_row_is_refused(cell): + with pytest.raises(ValueError, match=EPISODE_INPUT_COLUMN): + await _build_from({**_seed_row(_STORED), EPISODE_INPUT_COLUMN: cell}) + + +async def test_a_row_without_a_carrier_is_built_as_before(): + built = await _build_from({"persona": dict(_STORED["persona"]), "probe_type": "general_open_ended", "theme": "x"}) + + assert built["theme"] == "x" diff --git a/tests/test_cli_simulate_inputs.py b/tests/test_cli_simulate_inputs.py new file mode 100644 index 0000000..a353c39 --- /dev/null +++ b/tests/test_cli_simulate_inputs.py @@ -0,0 +1,863 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``usersim simulate --inputs``: stored rows run exactly as stored, one run per setup. + +The invariants: + +- two setups run over the same stored rows meet the same simulated users, in + the same scenario, and pair on ``input_id``; +- each setup gets its own run and its own ``trajectory_id`` per row, so a + resume never mixes two setups and never skips a row it did not run; +- everything that can change a conversation is part of the setup, and + nothing that cannot is. + +These run the real Data Designer pipeline. Every simulator-side model call is +stubbed, health checks are skipped, and the provider points at a closed local +port, so a call that slips past the stubs fails at once instead of going out. +""" + +from __future__ import annotations + +import copy +import importlib +import json +import logging +import shutil +import sys +from contextlib import ExitStack, contextmanager +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch + +import numpy as np +import pytest +from test_simulator_pipeline_e2e import _mock_call_llm_side_effect + +from usersim import cli +from usersim.cli import _inputs +from usersim.cli._inputs import ( + INPUT_ID_COLUMN, + load_episode_inputs, + prepare_episode_inputs, + setup_fingerprint, + simulate_inputs_to_dataset, +) +from usersim.cli._models import ModelsConfig, ProviderSpec, default_models_path, load_models_config +from usersim.cli._models import to_model_configs as _real_to_model_configs +from usersim.cli._pipeline import materialize_episode_inputs +from usersim.engine import generator as generator_module +from usersim.engine.core._assets import packaged_assets_dir +from usersim.engine.core.manifest import load_run_manifest +from usersim.engine.core.missing import none_for_missing +from usersim.engine.core.storage import list_runs, read_partitioned_dataset + +#: Every simulator-side module that holds its own reference to ``acall_llm``, +#: plus the module itself for the code that imports it at call time. +_LLM_MODULES = ( + "usersim.engine.core.llm", + "usersim.engine.core.simulation", + "usersim.engine.core.judges", + "usersim.engine.core.context", + "usersim.engine.core.translation", + "usersim.engine.core.realized_audit", + "usersim.engine.probes.health_disclosure.move_runtime", + "usersim.engine.probes.safety_agentic.generator", + "usersim.engine.probes.safety_chat_pressure.classifier", + "usersim.engine.probes.tool_calling.generator", +) +_LLM_SITES = tuple(m for m in _LLM_MODULES if hasattr(importlib.import_module(m), "acall_llm")) + + +def _probe_variant_pairs() -> list[tuple[str, str]]: + """Every registered probe paired with every variant it declares.""" + from usersim.engine.core.probes import known_probes, resolve_probe + + pairs = [] + for probe in sorted(known_probes()): + module = sys.modules[resolve_probe(probe).__module__] + pairs.extend((probe, str(v)) for v in module.PROBE_VARIANTS) + return pairs + + +#: A closed port: anything that reaches a provider fails at once. +_UNREACHABLE = "http://127.0.0.1:9/v1" + + +def _models(provider_api_key: str = "NVIDIA_API_KEY", **by_alias: dict[str, Any]) -> ModelsConfig: + """The default models, served from an unreachable provider, with per-alias changes.""" + base = load_models_config(default_models_path()) + provider = ProviderSpec(name="nvidia", endpoint=_UNREACHABLE, provider_type="openai", api_key=provider_api_key) + specs = tuple(replace(spec, **by_alias.get(spec.alias, {})) for spec in base.models) + return ModelsConfig(providers=(provider,), models=specs) + + +@contextmanager +def _offline(): + """Stub every simulator-side model call and skip Data Designer's health checks.""" + + def no_health_checks(models): + return [m.model_copy(update={"skip_health_check": True}) for m in _real_to_model_configs(models)] + + with ExitStack() as stack: + stack.enter_context(patch("usersim.cli._pipeline.to_model_configs", no_health_checks)) + for site in _LLM_SITES: + stack.enter_context(patch(f"{site}.acall_llm", side_effect=_mock_call_llm_side_effect)) + yield + + +@pytest.fixture(autouse=True) +def _environment(monkeypatch, tmp_path): + # A Data Designer preview, which materializing rows runs, leaves an + # artifacts folder in the working directory. + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("NVIDIA_API_KEY", "not-used") + # A value left set in the shell would change what the health variants run. + monkeypatch.delenv("USERSIM_DISCLOSURE_MOVES", raising=False) + + +def _materialize(probe_mix: dict[str, float], num_rows: int, seed: int = 1, max_turns: int = 1) -> list[dict[str, Any]]: + return materialize_episode_inputs( + locale="en_US", + num_rows=num_rows, + models=_models(), + assets_dir=packaged_assets_dir(), + probe_mix=probe_mix, + random_seed=seed, + max_turns=max_turns, + ) + + +@pytest.fixture(scope="module") +def _stored_rows(tmp_path_factory) -> list[dict[str, Any]]: + """Four stored rows from one materialization, with both probes among them, made once per module. + + The probe of each row is drawn at random, and the offline persona sample + holds only four people, so draw again until both probes appear. + """ + with pytest.MonkeyPatch.context() as env: + # What ``_environment`` sets for each test, for a fixture that outlives one. + env.chdir(tmp_path_factory.mktemp("materialize")) + env.setenv("NVIDIA_API_KEY", "not-used") + env.delenv("USERSIM_DISCLOSURE_MOVES", raising=False) + for _ in range(20): + stored = _materialize({"general_open_ended": 0.5, "tool_calling": 0.5}, num_rows=4) + if {row["probe_type"] for row in stored} == {"general_open_ended", "tool_calling"}: + return stored + raise AssertionError("twenty draws never mixed the two probes") + + +@pytest.fixture +def rows(_stored_rows) -> list[dict[str, Any]]: + """The module's stored rows, copied, so a test can change its own.""" + return copy.deepcopy(_stored_rows) + + +def _run(rows, out, models=None, *, assets_dir=None, **kwargs): + with _offline(): + inputs = prepare_episode_inputs(rows, assets_dir=assets_dir) + return simulate_inputs_to_dataset(inputs=inputs, models=models or _models(), out=out, **kwargs) + + +def _saved(out: Path, run_id: str): + return read_partitioned_dataset(out, run=run_id) + + +def _fingerprint(rows, models=None, *, assets_dir=None) -> str: + return setup_fingerprint(models or _models(), prepare_episode_inputs(rows, assets_dir=assets_dir))[0] + + +# ── two setups, one set of stored rows ────────────────────────────────── + + +def test_two_setups_get_their_own_runs_and_pair_on_input_id(rows, tmp_path): + first = _run(rows, tmp_path) + second = _run(rows, tmp_path, _models(assistant_model={"extra_body": {"reasoning_effort": "low"}})) + + assert first.run_id != second.run_id + assert (first.written, second.written) == (len(rows), len(rows)) + x, y = _saved(tmp_path, first.run_id), _saved(tmp_path, second.run_id) + assert set(x["trajectory_id"]).isdisjoint(y["trajectory_id"]) + pairs = x.merge(y, on=INPUT_ID_COLUMN, suffixes=("_x", "_y")) + assert sorted(pairs[INPUT_ID_COLUMN]) == sorted(row["trajectory_id"] for row in rows) + + +def test_both_setups_offer_each_tool_calling_row_the_same_tools(rows, tmp_path): + """The conversation runs with the stored id, from which ``tool_calling`` draws its tools.""" + first = _run(rows, tmp_path) + second = _run(rows, tmp_path, _models(assistant_model={"temperature": 0.1})) + + pairs = _saved(tmp_path, first.run_id).merge(_saved(tmp_path, second.run_id), on=INPUT_ID_COLUMN) + tools = pairs[pairs["probe_type_x"] == "tool_calling"] + assert len(tools) + assert tools["tool_subset_x"].astype(str).tolist() == tools["tool_subset_y"].astype(str).tolist() + + +def test_a_resume_reopens_the_run_with_the_same_setup(rows, tmp_path): + first = _run(rows, tmp_path) + _run(rows, tmp_path, _models(assistant_model={"temperature": 0.1})) + + again = _run(rows, tmp_path) + + assert (again.run_id, again.resumed, again.written, again.skipped) == (first.run_id, True, 0, len(rows)) + + +def test_without_skip_existing_each_call_starts_a_run(rows, tmp_path): + first = _run(rows, tmp_path, skip_existing=False) + second = _run(rows, tmp_path, skip_existing=False) + + assert first.run_id != second.run_id + assert len(list_runs(tmp_path)) == 2 + + +# ── what is part of a setup ───────────────────────────────────────────── + + +@pytest.mark.parametrize( + "models", + [ + pytest.param(lambda: _models(provider_api_key="ANOTHER_KEY_VARIABLE"), id="api_key"), + pytest.param(lambda: _models(assistant_model={"max_parallel_requests": 32}), id="max_parallel_requests"), + ], +) +def test_a_setting_that_cannot_change_a_conversation_keeps_the_setup(rows, models): + assert _fingerprint(rows, models()) == _fingerprint(rows) + + +def test_debug_logging_keeps_the_setup(rows, monkeypatch): + before = _fingerprint(rows) + monkeypatch.setenv("USERSIM_DEBUG_LOG", "1") + + assert _fingerprint(rows) == before + + +@pytest.mark.parametrize( + "models", + [ + pytest.param(lambda: _models(assistant_model={"timeout": 5.0}), id="timeout"), + pytest.param(lambda: _models(assistant_model={"temperature": 0.1}), id="temperature"), + pytest.param(lambda: _models(judge_model={"extra_body": {"reasoning_effort": "low"}}), id="judge"), + pytest.param(lambda: _models(assistant_model={"model": "another/model"}), id="model"), + ], +) +def test_a_model_setting_that_can_change_a_conversation_is_a_new_setup(rows, models): + assert _fingerprint(rows, models()) != _fingerprint(rows) + + +def test_a_providers_own_request_settings_are_part_of_the_setup(rows, monkeypatch): + """Data Designer merges a provider's extra_body and extra_headers into every request it serves.""" + from data_designer.config.models import ModelProvider + + models = ModelsConfig(providers=(), models=load_models_config(default_models_path()).models) + + def providers(**extra): + provider = ModelProvider(name="nvidia", endpoint=_UNREACHABLE, provider_type="openai", api_key="K", **extra) + monkeypatch.setattr("data_designer.config.default_model_settings.get_default_providers", lambda: [provider]) + + providers() + plain = setup_fingerprint(models, prepare_episode_inputs(rows))[0] + providers(extra_body={"reasoning_effort": "low"}) + with_body = setup_fingerprint(models, prepare_episode_inputs(rows))[0] + providers(extra_headers={"Authorization": "Bearer secret-value"}) + with_header, record = setup_fingerprint(models, prepare_episode_inputs(rows)) + + assert len({plain, with_body, with_header}) == 3 + assert "secret-value" not in json.dumps(record) + + +def test_changed_assets_are_a_new_setup_and_moved_ones_are_not(rows, tmp_path): + first, second = tmp_path / "assets", tmp_path / "elsewhere" + shutil.copytree(packaged_assets_dir(), first) + shutil.copytree(packaged_assets_dir(), second) + moved = _fingerprint(rows, assets_dir=first) + readme = next(first.rglob("README.md")) + readme.write_text(readme.read_text() + "\nedited\n") + + assert _fingerprint(rows, assets_dir=second) == moved + assert _fingerprint(rows, assets_dir=first) != moved + + +def test_a_change_in_a_lower_priority_asset_root_is_a_new_setup(monkeypatch, tmp_path): + """A spec that ``extends`` its own probe reads the copy below it, so every root counts.""" + top, below = tmp_path / "top", tmp_path / "below" + shutil.copytree(packaged_assets_dir(), top) + shutil.copytree(packaged_assets_dir(), below) + monkeypatch.setenv("USERSIM_ASSET_PATH", str(below)) + before = _inputs._asset_roots_digest(top) + shadowed = next(below.rglob("spec.yaml")) + shadowed.write_text(shadowed.read_text() + "\n# edited\n") + + assert _inputs._asset_roots_digest(top) != before + + +def test_a_data_designer_release_is_part_of_the_setup(rows, monkeypatch): + before = _fingerprint(rows) + monkeypatch.setattr(_inputs, "_data_designer_version", lambda: "999.0") + + assert _fingerprint(rows) != before + + +def test_the_setup_is_known_on_a_machine_where_data_designer_never_ran(rows, monkeypatch): + """Without its providers file, Data Designer serves the built-in providers it would write there.""" + + def missing(): + raise FileNotFoundError("model_providers.yaml") + + models = ModelsConfig(providers=(), models=load_models_config(default_models_path()).models) + monkeypatch.setattr("data_designer.config.default_model_settings.get_default_providers", missing) + + fingerprint, record = setup_fingerprint(models, prepare_episode_inputs(rows)) + assert record["models"]["assistant_model"]["provider"]["endpoint"] + + +def test_a_prompt_version_is_part_of_the_setup(rows, monkeypatch): + from usersim.engine.core.probes import resolve_probe + + before = _fingerprint(rows) + module = sys.modules[resolve_probe("general_open_ended").__module__] + monkeypatch.setattr(module, "PROMPT_VERSION", "v999", raising=False) + + assert _fingerprint(rows) != before + + +def test_a_behaviour_switch_in_the_environment_is_a_new_setup(rows, monkeypatch): + before = _fingerprint(rows) + monkeypatch.setenv("USERSIM_DISCLOSURE_STANDIN", "1") + + assert _fingerprint(rows) != before + + +def test_a_bank_override_counts_by_its_contents(rows, monkeypatch, tmp_path): + bank = tmp_path / "bank.json" + bank.write_text('{"items": [1]}') + elsewhere = tmp_path / "elsewhere.json" + elsewhere.write_text('{"items": [1]}') + monkeypatch.setenv("USERSIM_SAFETY_CHAT_PRESSURE_BANK", str(bank)) + original = _fingerprint(rows) + monkeypatch.setenv("USERSIM_SAFETY_CHAT_PRESSURE_BANK", str(elsewhere)) + moved = _fingerprint(rows) + elsewhere.write_text('{"items": [2]}') + + assert moved == original + assert _fingerprint(rows) != original + + +def test_an_extension_release_is_part_of_the_setup(rows, monkeypatch): + # The suite runs with discovery off; this test is about what it finds. + monkeypatch.delenv("USERSIM_DISABLE_EXTENSIONS", raising=False) + + def installed(version: str): + entry = SimpleNamespace( + name="partner_probe", + value="partner.probes:PartnerProbe", + dist=SimpleNamespace(name="partner", version=version), + ) + monkeypatch.setattr( + "importlib.metadata.entry_points", lambda group: [entry] if group == "usersim.probes" else [] + ) + + installed("1.0") + first = _fingerprint(rows) + installed("2.0") + + assert _fingerprint(rows) != first + + +# ── changed rows, concurrent writers, failures ────────────────────────── + + +def test_a_changed_stored_row_is_refused_before_any_model_call(rows, tmp_path): + first = _run(rows, tmp_path) + changed = copy.deepcopy(rows) + changed[0]["theme"] = "a theme edited after the run" + + with ( + patch("usersim.cli._pipeline.build_inputs_config_builder", side_effect=AssertionError("simulated")), + pytest.raises(ValueError, match="different contents"), + ): + _run(changed, tmp_path) + assert len(_saved(tmp_path, first.run_id)) == len(rows) + + +def _stale_plan(monkeypatch): + """Make a run plan as if it had started before anything was saved.""" + plan = _inputs.plan_inputs_run + + def before_anything_was_saved(inputs, models, out, *, skip_existing=True): + fresh = plan(inputs, models, out, skip_existing=False) + return replace(fresh, matching_run=None, pending={row["trajectory_id"] for row in inputs.rows}, saved=0) + + monkeypatch.setattr(_inputs, "plan_inputs_run", before_anything_was_saved) + + +def test_rows_another_writer_saved_meanwhile_are_not_saved_twice(rows, tmp_path, monkeypatch): + first = _run(rows, tmp_path) + _stale_plan(monkeypatch) + + second = _run(rows, tmp_path) + + assert (second.run_id, second.written, second.skipped) == (first.run_id, 0, len(rows)) + assert list_runs(tmp_path) == [first.run_id] + assert len(_saved(tmp_path, first.run_id)) == len(rows) + + +def test_a_conflict_saved_meanwhile_is_refused_at_the_write(rows, tmp_path, monkeypatch): + first = _run(rows, tmp_path) + changed = copy.deepcopy(rows) + changed[0]["theme"] = "a theme edited after the run" + _stale_plan(monkeypatch) + + with pytest.raises(ValueError, match="different contents"): + _run(changed, tmp_path) + assert len(_saved(tmp_path, first.run_id)) == len(rows) + + +def test_a_failure_before_any_save_leaves_no_run(rows, tmp_path): + with ( + patch("usersim.cli._pipeline.build_inputs_config_builder", side_effect=RuntimeError("simulated")), + pytest.raises(RuntimeError), + ): + _run(rows, tmp_path) + + assert list_runs(tmp_path) == [] + assert not list(tmp_path.glob(".claim-*")) + + +# ── the stored rows reach the probe, and the saved rows look like a plain run's ── + + +@contextmanager +def _built_rows(): + """Collect the rows the probe is built from.""" + built: list[dict[str, Any]] = [] + construct = generator_module.construct_probe_episode + + def spy(data, **kwargs): + built.append(dict(data)) + return construct(data, **kwargs) + + with patch.object(generator_module, "construct_probe_episode", spy): + yield built + + +@contextmanager +def _undeclared_column_warnings(): + seen: list[str] = [] + + class _Capture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.levelno >= logging.WARNING and "side_effect_columns" in record.getMessage(): + seen.append(record.getMessage()) + + handler, engine_logger = _Capture(), logging.getLogger("usersim.engine") + engine_logger.addHandler(handler) + try: + yield seen + finally: + engine_logger.removeHandler(handler) + + +@pytest.mark.parametrize(("probe_type", "variant"), _probe_variant_pairs()) +def test_every_probe_is_built_from_its_stored_row_exactly(probe_type, variant, tmp_path): + (stored,) = _materialize({probe_type: 1.0}, num_rows=1) + stored["probe_variant"] = variant + + with _built_rows() as built, _undeclared_column_warnings() as warnings: + result = _run([stored], tmp_path) + + (row,) = built + expected = none_for_missing({key: value for key, value in stored.items() if key != "usersim_config"}) + # The probe gets plain Python values, while a stored row may hold numpy + # scalars (a pandas 2 preview leaves `theme` as one). + for key, value in expected.items(): + want = value.item() if isinstance(value, np.generic) else value + assert row[key] == want and type(row[key]) is type(want), key + assert warnings == [] + (saved_id,) = _saved(tmp_path, result.run_id)[INPUT_ID_COLUMN] + assert saved_id == stored["trajectory_id"] + + +def test_values_keep_their_types_across_rows(tmp_path): + """A seed table would turn this int into a float, because another row holds a float.""" + stored = _materialize({"general_open_ended": 1.0}, num_rows=2) + stored[0]["extension_score"], stored[1]["extension_score"] = 9007199254740993, 0.5 + stored[0]["extension_memory"] = {"sessions": [{"week": 1, "notes": ["slept badly"]}], "flag": True} + stored[1]["extension_memory"] = {"sessions": [], "flag": False} + + with _built_rows() as built: + _run(stored, tmp_path) + + by_id = {row["trajectory_id"]: row for row in built} + for row in stored: + for key in ("extension_score", "extension_memory"): + got = by_id[row["trajectory_id"]][key] + assert got == row[key] and type(got) is type(row[key]) + + +def test_an_override_column_on_a_stored_row_reaches_the_probe(tmp_path): + """A probe's override column, set on a stored row, picks what the probe runs. + + ``usersim simulate --panel`` reads only a panel's locales, so this is the + way to pin an override: materialize rows, set the column, run them. + """ + import yaml + + (stored,) = _materialize({"safety_chat_pressure": 1.0}, num_rows=1) + first = _run([stored], tmp_path / "derived") + (derived,) = _saved(tmp_path / "derived", first.run_id)["strategy_id"] + bank = yaml.safe_load((packaged_assets_dir() / "safety_chat_pressure" / "sample.yaml").read_text(encoding="utf-8")) + chosen = next(strategy["id"] for strategy in bank["strategies"] if strategy["id"] != derived) + + stored["pressure_strategy_id"] = chosen + second = _run([stored], tmp_path / "chosen") + + assert list(_saved(tmp_path / "chosen", second.run_id)["strategy_id"]) == [chosen] + + +@pytest.mark.parametrize("probe_type", ["general_open_ended", "tool_calling"]) +def test_saved_rows_have_a_plain_runs_columns_plus_input_id(probe_type, tmp_path): + """The columns of a plain run, and the stored values in them as stored. + + The ``persona_*`` columns, and ``tool_calling``'s toolset columns, are + added back from the stored rows, so they must keep the stored text and + numbers rather than come back encoded. + """ + from usersim.cli.simulate import _simulate_to_parquet + + plain_out, inputs_out = tmp_path / "plain", tmp_path / "inputs" + with _offline(): + _simulate_to_parquet( + models=_models(), + models_path=None, + panel=None, + locales=["en_US"], + num_rows=2, + out=plain_out, + assets_dir=packaged_assets_dir(), + max_turns=1, + probe_mix={probe_type: 1.0}, + random_seed=1, + skip_existing=True, + ) + stored = _materialize({probe_type: 1.0}, num_rows=2) + result = _run(stored, inputs_out) + + plain = read_partitioned_dataset(plain_out) + saved = _saved(inputs_out, result.run_id) + assert set(saved.columns) == set(plain.columns) | {INPUT_ID_COLUMN} + by_input = saved.set_index(INPUT_ID_COLUMN) + for row in stored: + shared = sorted(set(row) & set(plain.columns) - {"trajectory_id"}) + values = none_for_missing({key: by_input.loc[row["trajectory_id"], key] for key in shared}) + for key in shared: + got, want = (value.item() if isinstance(value, np.generic) else value for value in (values[key], row[key])) + assert got == want and type(got) is type(want), key + + +@pytest.mark.parametrize( + ("probe_type", "variant"), + [("general_open_ended", "default"), ("tool_calling", "default"), ("health_therapy_disclosure", "guarded")], +) +def test_only_input_id_holds_the_stored_id(probe_type, variant, tmp_path): + """No other saved value contains a stored id, JSON columns included. + + The conversation runs under the stored id, and two of these probes read it + mid-run: ``tool_calling`` draws its tools from it, and a guarded health + variant stamps it on every move it reports to a move hook, which therefore + sees the stored id, the saved row's ``input_id``. Moves start on the second + turn. + """ + from usersim.engine.probes.health_disclosure.mixin import GuardedMoveMixin + + stored = _materialize({probe_type: 1.0}, num_rows=2, max_turns=3) + for row in stored: + row["probe_variant"] = variant + reported: list[dict[str, Any]] = [] + previous = GuardedMoveMixin.register_move_hook(reported.append) + try: + result = _run(stored, tmp_path) + finally: + GuardedMoveMixin.register_move_hook(previous) + + saved = _saved(tmp_path, result.run_id) + stored_ids = {row["trajectory_id"] for row in stored} + assert set(saved[INPUT_ID_COLUMN]) == stored_ids + if variant == "guarded": + assert {payload["trajectory_id"] for payload in reported} == stored_ids + for column in saved.columns.drop(INPUT_ID_COLUMN): + text = saved[column].astype(str) + for row in stored: + assert not text.str.contains(row["trajectory_id"], regex=False).any(), column + + +def test_the_manifest_describes_the_run(rows, tmp_path, monkeypatch): + source = tmp_path / "rows.jsonl" + source.write_text("".join(json.dumps(row) + "\n" for row in rows)) + # The invocation started long before its first rows were saved. + monkeypatch.setattr(_inputs, "time", SimpleNamespace(time=lambda: 1_700_000_000.4)) + with _offline(): + inputs = prepare_episode_inputs(load_episode_inputs(source), source=source) + result = simulate_inputs_to_dataset(inputs=inputs, models=_models(), out=tmp_path / "out") + + manifest = load_run_manifest(result.run_id, tmp_path / "out") + assert int(result.run_id) > 1_700_000_000 + assert manifest.started_at_epoch == 1_700_000_000 + assert manifest.scope.panel_path == str(source.resolve()) + assert manifest.scope.panel_content_hash + assert manifest.scope.trajectories_requested == {"en_US": len(rows)} + assert manifest.simulation_config.max_turns == rows[0]["usersim_config"]["max_turns"] + assert manifest.source.kind == "cli" and manifest.source.invocation + + +def test_a_plain_run_records_when_it_started_and_leaves_no_artifacts(tmp_path, monkeypatch): + """Data Designer names artifacts by the second a batch starts; each batch gets its own folder.""" + from usersim.cli import simulate + + monkeypatch.setattr(simulate, "time", SimpleNamespace(time=lambda: 1_700_000_000.4)) + with _offline(): + simulate._simulate_to_parquet( + models=_models(), + models_path=None, + panel=None, + locales=["en_US"], + num_rows=2, + out=tmp_path / "plain", + assets_dir=packaged_assets_dir(), + max_turns=1, + probe_mix={"general_open_ended": 1.0}, + random_seed=1, + skip_existing=True, + ) + + (run_id,) = list_runs(tmp_path / "plain") + assert load_run_manifest(run_id, tmp_path / "plain").started_at_epoch == 1_700_000_000 + assert not (tmp_path / "artifacts").exists() + + +def test_an_inputs_run_leaves_no_artifacts(rows, tmp_path): + """Its Data Designer batches write into temporary folders, not the working directory.""" + _run(rows, tmp_path / "out") + + assert not (tmp_path / "artifacts").exists() + + +def test_saved_columns_keep_their_stored_values(tmp_path): + """A column of one plain type is saved as it is; anything Arrow would change is saved as JSON text.""" + stored = _materialize({"general_open_ended": 1.0}, num_rows=2) + stored[0]["extension_count"], stored[1]["extension_count"] = 9007199254740993, None + stored[0]["extension_score"], stored[1]["extension_score"] = 9007199254740993, 0.5 + stored[0]["extension_note"], stored[1]["extension_note"] = {"mood": "low"}, "plain text" + + result = _run(stored, tmp_path) + + saved = _saved(tmp_path, result.run_id).set_index(INPUT_ID_COLUMN) + first, second = (saved.loc[row["trajectory_id"]] for row in stored) + assert int(first["extension_count"]) == 9007199254740993 + assert [json.loads(first["extension_score"]), json.loads(second["extension_score"])] == [9007199254740993, 0.5] + assert [json.loads(first["extension_note"]), json.loads(second["extension_note"])] == [ + {"mood": "low"}, + "plain text", + ] + + +def _two_rows_with(first: Any, second: Any) -> list[dict[str, Any]]: + stored = _materialize({"general_open_ended": 1.0}, num_rows=2) + stored[0]["extension_value"], stored[1]["extension_value"] = first, second + return stored + + +def test_a_resume_saves_each_column_as_the_first_invocation_did(tmp_path): + """Whether a column is saved as JSON text is chosen from all the stored rows, not the ones left to run.""" + stored = _two_rows_with(9007199254740993, 0.5) + with _offline(): + inputs = prepare_episode_inputs(stored) + plan = _inputs.plan_inputs_run(inputs, _models(), tmp_path) + # The first invocation stops after saving the first row. + first = simulate_inputs_to_dataset( + inputs=inputs, models=_models(), out=tmp_path, plan=replace(plan, pending={stored[0]["trajectory_id"]}) + ) + second = _run(stored, tmp_path) + + assert second.run_id == first.run_id and (first.written, second.written) == (1, 1) + saved = _saved(tmp_path, first.run_id).set_index(INPUT_ID_COLUMN)["extension_value"] + assert [json.loads(saved[row["trajectory_id"]]) for row in stored] == [9007199254740993, 0.5] + + +def test_batches_whose_fields_keep_their_type_read_back_exactly(tmp_path): + """Rows run in two calls into one run, a field holding one type in both.""" + stored = _two_rows_with(9007199254740993, 7) + first = _run(stored[:1], tmp_path) + second = _run(stored[1:], tmp_path) + + assert second.run_id == first.run_id + saved = _saved(tmp_path, first.run_id).set_index(INPUT_ID_COLUMN)["extension_value"] + assert [int(saved[row["trajectory_id"]]) for row in stored] == [9007199254740993, 7] + + +def test_concurrent_invocations_save_each_row_once(rows, tmp_path): + import threading + + results, errors = [], [] + + def invoke(): + try: + inputs = prepare_episode_inputs(rows) + results.append(simulate_inputs_to_dataset(inputs=inputs, models=_models(), out=tmp_path / "out")) + except Exception as error: # surfaced below + errors.append(error) + + with _offline(): + threads = [threading.Thread(target=invoke) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=300) + + assert errors == [] + (run_id,) = list_runs(tmp_path / "out") + saved = _saved(tmp_path / "out", run_id) + assert sorted(saved[INPUT_ID_COLUMN]) == sorted(row["trajectory_id"] for row in rows) + assert sum(result.written for result in results) == len(rows) + + +def test_rows_stay_in_the_run_the_plan_found(rows, tmp_path, monkeypatch): + """A newer run with the same setup, made meanwhile, does not split the rows across two runs.""" + found = _run(rows[:2], tmp_path) + newer = _run(rows[:2], tmp_path, skip_existing=False) + matching = _inputs._matching_run + calls = [] + + def first_the_found_run_then_the_newer_one(out, fingerprint): + calls.append(fingerprint) + return found.run_id if len(calls) == 1 else matching(out, fingerprint) + + monkeypatch.setattr(_inputs, "_matching_run", first_the_found_run_then_the_newer_one) + result = _run(rows, tmp_path) + + assert (result.run_id, result.written, result.skipped) == (found.run_id, len(rows) - 2, 2) + assert len(_saved(tmp_path, found.run_id)) == len(rows) + assert len(_saved(tmp_path, newer.run_id)) == 2 + + +def test_a_conflict_found_when_saving_exits_with_a_message(rows, tmp_path, monkeypatch): + out = tmp_path / "o" + argv = ["simulate", "--inputs", str(_write(rows, tmp_path / "rows.jsonl")), "--out", str(out)] + with _offline(): + assert cli.main(argv) == 0 + changed = _edited(rows, lambda r: r[0].update(theme="a theme edited after the run")) + argv[2] = str(_write(changed, tmp_path / "changed.jsonl")) + _stale_plan(monkeypatch) + with pytest.raises(SystemExit, match="different contents"): + cli.main(argv) + + +# ── checks before any model call ──────────────────────────────────────── + + +def _edited(rows, edit): + changed = copy.deepcopy(rows) + edit(changed) + return changed + + +@pytest.mark.parametrize( + ("edit", "message"), + [ + (lambda r: r[1]["usersim_config"].update(max_turns=9), "different settings"), + (lambda r: r[1].update(trajectory_id=r[0]["trajectory_id"]), "more than one row"), + (lambda r: r[0].update(input_id="x"), "input_id"), + (lambda r: r[0].update(usersim_episode_input="{}"), "usersim_episode_input"), + (lambda r: r[0].pop("usersim_config"), "usersim_config"), + (lambda r: r[0].update(probe_type="no_such_probe"), "no_such_probe"), + (lambda r: r[0].update(locale="pt_BR"), "locale"), + (lambda r: r[0].update(persona=None), "persona"), + (lambda r: r[0].update(extension_blob=b"bytes"), "JSON"), + (lambda r: r[0].update(trajectory_id=""), "trajectory_id"), + ], +) +def test_rows_that_cannot_run_are_refused(rows, edit, message): + with pytest.raises(ValueError, match=message): + prepare_episode_inputs(_edited(rows, edit)) + + +def test_a_missing_stored_asset_folder_is_an_error_until_one_is_given(rows): + moved = _edited(rows, lambda r: [row["usersim_config"].update(assets_dir="/no/such/assets") for row in r]) + + with pytest.raises(ValueError, match="--assets-dir"): + prepare_episode_inputs(moved) + assert prepare_episode_inputs(moved, assets_dir=packaged_assets_dir()).assets_dir == packaged_assets_dir() + + +def _write(rows, path: Path) -> Path: + path.write_text("".join(json.dumps(row) + "\n" for row in rows)) + return path + + +@pytest.mark.parametrize( + ("flag", "value"), + [ + ("--max-turns", "9"), + ("--probe-mix", "general_open_ended=1"), + ("--random-seed", "4"), + ("--no-store-reasoning", None), + ], +) +def test_a_flag_the_stored_rows_settle_is_refused(rows, tmp_path, flag, value): + argv = ["simulate", "--inputs", str(_write(rows, tmp_path / "rows.jsonl")), "--out", str(tmp_path / "o")] + with pytest.raises(SystemExit, match=flag): + cli.main([*argv, flag, *([value] if value else []), "--dry-run"]) + + +def test_a_dry_run_shows_the_plan_without_running(rows, tmp_path, capsys): + argv = [ + "simulate", + "--inputs", + str(_write(rows, tmp_path / "rows.jsonl")), + "--out", + str(tmp_path / "o"), + "--dry-run", + ] + with patch("usersim.cli._pipeline.build_inputs_config_builder", side_effect=AssertionError("ran")): + assert cli.main(argv) == 0 + + out = capsys.readouterr().out + assert "dry run" in out and "new, assigned when the first rows are saved" in out + assert f"max_turns={rows[0]['usersim_config']['max_turns']!r}" in out + assert list_runs(tmp_path / "o") == [] + + +def test_a_root_without_file_locks_still_gets_its_run_with_one_warning(rows, tmp_path, monkeypatch, caplog): + """Shared filesystems such as Lustre are often mounted without ``flock``; stored rows run there anyway.""" + fcntl = pytest.importorskip("fcntl") + import errno + + def unsupported(*args, **kwargs): + raise OSError(errno.ENOSYS, "flock is not supported here") + + monkeypatch.setattr(fcntl, "flock", unsupported) + with caplog.at_level(logging.WARNING, logger="usersim.engine"): + first = _run(rows[:2], tmp_path) + second = _run(rows, tmp_path) + + assert second.run_id == first.run_id and second.resumed + assert set(_saved(tmp_path, first.run_id)[INPUT_ID_COLUMN]) == {row["trajectory_id"] for row in rows} + assert sum("not supported" in record.getMessage() for record in caplog.records) == 1 + + +def test_rows_stored_as_parquet_load_as_they_were_written(rows, tmp_path): + """``--materialize-inputs`` writes parquet through pandas; reading it back must not re-type it.""" + import pandas as pd + + path = tmp_path / "rows.parquet" + pd.DataFrame(rows).to_parquet(path, index=False) + + loaded = load_episode_inputs(path) + inputs = prepare_episode_inputs(loaded) + + assert [row["trajectory_id"] for row in inputs.rows] == [row["trajectory_id"] for row in rows] + assert [row["usersim_config"]["max_turns"] for row in loaded] == [ + row["usersim_config"]["max_turns"] for row in rows + ] + assert all(isinstance(row["persona"], dict) for row in loaded)