Skip to content

fix(storage): write simulation runs under a lock and publish them whole #34

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
isarasua merged 3 commits into from
Oct 13, 2026
+1,013 −113
Merged

fix(storage): write simulation runs under a lock and publish them whole #34

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f2989c5
fix(storage): write simulation runs under a lock and publish them whole
isarasua Oct 9, 2026
e2ec9d9
Merge branch 'main' into isarasua/fix/locked-run-writes
isarasua Oct 12, 2026
f4045c9
fix(storage): write where flock is unsupported, and keep live writers…
isarasua Oct 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 74 additions & 59 deletions notebooks/01_simulate.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@
"import logging\n",
"import os\n",
"import subprocess\n",
"import tempfile\n",
"import time\n",
"import warnings\n",
"from collections import Counter\n",
"from pathlib import Path\n",
Expand Down Expand Up @@ -101,14 +103,16 @@
" preview_summary,\n",
")\n",
"from usersim.engine.core.storage import (\n",
" LEGACY_RUN_ID,\n",
" existing_trajectory_ids,\n",
" latest_run_id,\n",
" list_partition_keys,\n",
" new_run_id,\n",
" read_partitioned_dataset,\n",
" resolve_run_or_raise,\n",
" run_started_at,\n",
" run_subroot,\n",
" write_locale_partition,\n",
" sampled_run_to_resume,\n",
" save_sampled_rows,\n",
")\n",
"\n",
"warnings.filterwarnings(\"ignore\", message=\"urllib3.*chardet.*charset_normalizer\")\n",
Expand Down Expand Up @@ -577,12 +581,12 @@
"source": [
"## 🔁 Resumability check + run-id resolution\n",
"\n",
"Every fresh simulation lands under its own `run=<unix-epoch>/` partition under `TRAJECTORY_PATH`. The reader defaults to \"latest run\" so downstream evaluation + reporting see the most recent results without configuration. Old runs stay on disk untouched, available for run-over-run comparison or as historical evidence.\n",
"Every fresh simulation lands under its own `run=<unix-epoch>/` partition under `TRAJECTORY_PATH`, made when its first rows are saved. The reader defaults to \"latest run\" so downstream evaluation + reporting see the most recent results without configuration. Old runs stay on disk untouched, available for run-over-run comparison or as historical evidence.\n",
"\n",
"Behaviour depends on `RESUME` (set in the configuration cell above):\n",
"\n",
"- **`RESUME = True`**: resume into the most-recent existing run. Skip trajectory_ids already present in that run and append new partitions to it. Useful when iterating on a probe-mix expansion or recovering from a crashed run.\n",
"- **`RESUME = False`** (default): mint a new run id (`int(time.time())`) and start fresh under `TRAJECTORY_PATH/run=<that-id>/`. Old runs are not deleted (they keep being readable via `read_partitioned_dataset(TRAJECTORY_PATH, run=\"<id>\")` or `run=\"all\"`) but the default-latest read won't surface them.\n",
"- **`RESUME = True`**: resume into the most-recent run of sampled rows (never a run of another kind). Skip trajectory_ids already present in that run and append new partitions to it. Useful when iterating on a probe-mix expansion or recovering from a crashed run.\n",
"- **`RESUME = False`** (default): start a new run under `TRAJECTORY_PATH/run=<id>/`, made when its first rows are saved, so a failure before then leaves no empty run. Old runs are not deleted (they keep being readable via `read_partitioned_dataset(TRAJECTORY_PATH, run=\"<id>\")` or `run=\"all\"`) but the default-latest read won't surface them.\n",
"\n",
"The simulator's identity scheme (`trajectory_id` is content-hashed from persona + probe + seed + model identities + prompt version) means resumption is byte-stable: re-running the same configuration against the same panel produces the same `trajectory_id`s and the resume path is a no-op cost."
]
Expand Down Expand Up @@ -619,20 +623,21 @@
" f\"missing ones above.\"\n",
" )\n",
"elif RESUME:\n",
" RUN_ID = latest_run_id(TRAJECTORY_PATH) or new_run_id()\n",
" already_done = existing_trajectory_ids(TRAJECTORY_PATH, run=RUN_ID)\n",
" # None when there is no run of sampled rows yet: one is made with the first rows saved.\n",
" RUN_ID = sampled_run_to_resume(TRAJECTORY_PATH)\n",
" already_done = existing_trajectory_ids(TRAJECTORY_PATH, run=RUN_ID) if RUN_ID else set()\n",
" print(\n",
" f\"RESUME=True: resuming into run={RUN_ID}; \"\n",
" f\"RESUME=True: resuming into run={RUN_ID or '(new)'}; \"\n",
" f\"{len(already_done)} trajectory_id(s) already present, will skip those.\"\n",
" )\n",
"else:\n",
" RUN_ID = new_run_id()\n",
" RUN_ID = None\n",
" already_done = set()\n",
" print(\n",
" f\"RESUME=False: minting new run={RUN_ID} under {TRAJECTORY_PATH}/.\"\n",
" f\"RESUME=False: a new run under {TRAJECTORY_PATH}/ is made when the first rows are saved.\"\n",
" )\n",
" existing = latest_run_id(TRAJECTORY_PATH)\n",
" if existing is not None and existing != RUN_ID:\n",
" if existing is not None:\n",
" print(\n",
" f\" (Previous run on disk: {existing}. Untouched. Read it back via \"\n",
" f\"read_partitioned_dataset(TRAJECTORY_PATH, run='{existing}').)\"\n",
Expand Down Expand Up @@ -661,6 +666,7 @@
"from usersim.cli._pipeline import DEFAULT_TOOL_CALLING_THEMES\n",
"\n",
"total_written = 0\n",
"STARTED_AT = time.time()\n",
"# The toolset seed is a tool_calling input. Passing it when the mix has no\n",
"# tool_calling probe is rejected as a contradictory override, so only send it\n",
"# when that probe is actually selected.\n",
Expand All @@ -684,59 +690,68 @@
" toolset_seed_path=_seed_path,\n",
" random_seed=RANDOM_SEED,\n",
" )\n",
" result = data_designer.create(builder, num_records=NUM_ROWS_PER_LOCALE)\n",
" df = result.load_dataset()\n",
" if already_done:\n",
" df = df[~df[\"trajectory_id\"].isin(already_done)]\n",
" if df.empty:\n",
" print(f\" (locale={locale}: no new trajectories, all already in dataset)\")\n",
" continue\n",
" # Data Designer names a batch's artifacts by the second it started, so two\n",
" # kernels in one directory could share them; each batch gets its own folder.\n",
" with tempfile.TemporaryDirectory(prefix=\"usersim_simulate_\") as _artifacts:\n",
" result = data_designer.create(builder, num_records=NUM_ROWS_PER_LOCALE, artifact_path=_artifacts)\n",
" df = result.load_dataset()\n",
" # Atomic per-locale write: this locale's partition lands on disk\n",
" # before the next locale starts. A crash mid-run leaves earlier\n",
" # locales' partitions intact and resumable.\n",
" write_locale_partition(df, TRAJECTORY_PATH, locale=locale, run_id=RUN_ID)\n",
" total_written += len(df)\n",
" print(f\" wrote {len(df)} new trajectories for {locale}\")\n",
" # locales' partitions intact and resumable. It holds the run lock, skips\n",
" # rows the run already has, and makes the run with its first rows.\n",
" RUN_ID, written = save_sampled_rows(\n",
" df, TRAJECTORY_PATH, locale=locale, run_id=RUN_ID, started_at=STARTED_AT\n",
" )\n",
" if not written:\n",
" print(f\" (locale={locale}: no new trajectories, all already in dataset)\")\n",
" continue\n",
" total_written += written\n",
" print(f\" wrote {written} new trajectories for {locale} into run={RUN_ID}\")\n",
"\n",
"print(f\"\\nWrote {total_written} new trajectories to {TRAJECTORY_PATH}\")\n",
"\n",
"# Write the run manifest. ``RUN_ID`` is the unix-epoch start time\n",
"# (see ``new_run_id``), so it doubles as ``started_at_epoch`` and the\n",
"# manifest's runtime is finished_at - RUN_ID. First-writer-wins:\n",
"# resuming into a run that already has a manifest is a noop.\n",
"sim_cfg_for_manifest = ConversationSimulatorConfig(\n",
" name=\"conversation_messages\",\n",
" locale=\"placeholder\",\n",
" max_turns=MAX_TURNS,\n",
" max_assistant_attempts=MAX_ASSISTANT_ATTEMPTS,\n",
" context_compression=True,\n",
" compression_window=1,\n",
" random_seed=RANDOM_SEED,\n",
")\n",
"# Describe the RUN, not just this pass. On a resumed run LOCALES holds only the\n",
"# locales this pass filled in, so recording it verbatim would understate a run that\n",
"# already has others on disk -- and the manifest is first-writer-wins, so there is no\n",
"# second chance to correct it. Union with what is actually there.\n",
"_run_locales = sorted(\n",
" set(LOCALES) | set(list_partition_keys(run_subroot(TRAJECTORY_PATH, RUN_ID)))\n",
")\n",
"manifest_path = write_simulator_manifest(\n",
" run_id=RUN_ID,\n",
" started_at_epoch=int(RUN_ID),\n",
" trajectory_root=TRAJECTORY_PATH,\n",
" models_config=MODELS,\n",
" models_config_path=MODELS_CONFIG_PATH,\n",
" sim_config=sim_cfg_for_manifest,\n",
" locales_requested=_run_locales,\n",
" probe_mix=PROBE_MIX,\n",
" tool_calling_themes=DEFAULT_TOOL_CALLING_THEMES,\n",
" toolset_seed_path=TOOLSET_SEED_PATH if TOOLSET_SEED_PATH.exists() else None,\n",
" assets_dir=ASSETS_DIR,\n",
" trajectories_requested={l: NUM_ROWS_PER_LOCALE for l in _run_locales},\n",
" source_kind=\"notebook\",\n",
" source_invocation=\"notebooks/01_simulate.ipynb\",\n",
")\n",
"print(f\"Wrote run manifest to {manifest_path}\")"
"# Write the run manifest. ``started_at_epoch`` is when the pass that made\n",
"# the run started, kept in the run's record. First-writer-wins: resuming\n",
"# into a run that already has a manifest is a noop. No rows saved means no\n",
"# run, and so no manifest. The bare legacy layout has no run folder to hold one.\n",
"if RUN_ID is None:\n",
" print(\"No rows were saved, so no run was made.\")\n",
"elif RUN_ID == LEGACY_RUN_ID:\n",
" print(\"This pass resumed the bare legacy layout, which has no run folder for a manifest.\")\n",
"else:\n",
" sim_cfg_for_manifest = ConversationSimulatorConfig(\n",
" name=\"conversation_messages\",\n",
" locale=\"placeholder\",\n",
" max_turns=MAX_TURNS,\n",
" max_assistant_attempts=MAX_ASSISTANT_ATTEMPTS,\n",
" context_compression=True,\n",
" compression_window=1,\n",
" random_seed=RANDOM_SEED,\n",
" )\n",
" # Describe the RUN, not just this pass. On a resumed run LOCALES holds only the\n",
" # locales this pass filled in, so recording it verbatim would understate a run that\n",
" # already has others on disk -- and the manifest is first-writer-wins, so there is no\n",
" # second chance to correct it. Union with what is actually there.\n",
" _run_locales = sorted(\n",
" set(LOCALES) | set(list_partition_keys(run_subroot(TRAJECTORY_PATH, RUN_ID)))\n",
" )\n",
" manifest_path = write_simulator_manifest(\n",
" run_id=RUN_ID,\n",
" started_at_epoch=run_started_at(TRAJECTORY_PATH, RUN_ID) or int(STARTED_AT),\n",
" trajectory_root=TRAJECTORY_PATH,\n",
" models_config=MODELS,\n",
" models_config_path=MODELS_CONFIG_PATH,\n",
" sim_config=sim_cfg_for_manifest,\n",
" locales_requested=_run_locales,\n",
" probe_mix=PROBE_MIX,\n",
" tool_calling_themes=DEFAULT_TOOL_CALLING_THEMES,\n",
" toolset_seed_path=TOOLSET_SEED_PATH if TOOLSET_SEED_PATH.exists() else None,\n",
" assets_dir=ASSETS_DIR,\n",
" trajectories_requested={l: NUM_ROWS_PER_LOCALE for l in _run_locales},\n",
" source_kind=\"notebook\",\n",
" source_invocation=\"notebooks/01_simulate.ipynb\",\n",
" )\n",
" print(f\"Wrote run manifest to {manifest_path}\")"
]
},
{
Expand Down
2 changes: 2 additions & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,11 @@ dependencies = [
# Reached directly rather than through another package's public API, so
# declared even though each also arrives transitively today: a transitive
# pin can be dropped by its parent without warning.
# filelock -- core/storage.py locks a trajectory root while writing
# httpx -- cli/_models.py probes model endpoints
# huggingface_hub -- selection/io.py publishes dataset cards
# numpy -- asset_gen near-duplicate detection
"filelock",
"httpx",
"huggingface-hub",
"numpy",
Expand Down
86 changes: 45 additions & 41 deletions src/usersim/cli/simulate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@

import argparse
import logging
import tempfile
import time
from fractions import Fraction
from pathlib import Path

Expand Down Expand Up @@ -466,35 +468,29 @@ def _simulate_to_parquet(
from usersim.engine.core.locale import persona_dataset_locale
from usersim.engine.core.manifest import write_simulator_manifest
from usersim.engine.core.storage import (
LEGACY_RUN_ID,
existing_trajectory_ids,
latest_run_id,
new_run_id,
write_locale_partition,
run_started_at,
sampled_run_to_resume,
save_sampled_rows,
)

started_at = time.time()
out = Path(out).resolve()
out.mkdir(parents=True, exist_ok=True)

# Resolve the target run id. ``--no-skip-existing`` means
# "fresh start" → mint a new run id. Default (``--skip-existing``)
# means "resume into the latest existing run". On a first
# invocation with no existing runs, both paths land on a fresh
# ``new_run_id()``.
if skip_existing:
run_id = latest_run_id(out) or new_run_id()
else:
run_id = new_run_id()
logger.info("run_id=%s (%s)", run_id, "resume" if skip_existing else "fresh")
print(f" run_id: {run_id}")

# Skip-existing scopes to the resolved run only (not all runs).
# Resuming into an existing run skips trajectory_ids already
# written for THIS run; trajectory_ids in other runs are
# independent (they had different model identities, prompts,
# seeds — content-hashed trajectory_ids would still collide if
# everything matched, but cross-run skipping is not the
# behaviour we want).
already_done = existing_trajectory_ids(out, run=run_id) if skip_existing else set()
# ``--no-skip-existing`` starts a fresh run; the default resumes the
# latest run this command made. A fresh run is claimed only when its
# first rows are saved, so a failure before then leaves no empty run.
run_id = sampled_run_to_resume(out) if skip_existing else None
logger.info("run_id=%s (%s)", run_id or "new", "resume" if run_id else "fresh")
print(f" run_id: {run_id or 'new, assigned when the first rows are saved'}")

# Rows already saved in the resumed run are skipped when rows are saved
# (``save_sampled_rows`` checks again under the lock); this count is only
# reported. Only this run counts: an id saved in another run is
# independent of this one.
already_done = existing_trajectory_ids(out, run=run_id) if run_id else set()
if already_done:
logger.info(
"run=%s already has %d trajectories — will skip",
Expand Down Expand Up @@ -563,23 +559,27 @@ def _simulate_to_parquet(
n = _count_rows_for_locale(panel, locale)
else:
n = num_rows
result = data_designer.create(builder, num_records=n)
df = result.load_dataset()
if already_done:
df = df[~df["trajectory_id"].isin(already_done)]
if df.empty:
logger.info("locale=%s: no new trajectories", locale)
continue
# Data Designer names a batch's artifacts by the second it started, so
# two runs in one directory could share them; each batch gets its own
# folder, removed once its rows are loaded.
with tempfile.TemporaryDirectory(prefix="usersim_simulate_") as artifacts:
df = data_designer.create(builder, num_records=n, artifact_path=artifacts).load_dataset()
# Per-locale atomic write: this locale's partition lands on
# disk before the next locale starts. A crash mid-run leaves
# earlier locales' partitions intact and resumable (within
# the same run id).
write_locale_partition(df, out, locale=locale, run_id=run_id)
total_written += len(df)
claimed = run_id is None
run_id, written = save_sampled_rows(df, out, locale=locale, run_id=run_id, started_at=started_at)
if not written:
logger.info("locale=%s: no new trajectories", locale)
continue
if claimed:
print(f" run_id: {run_id}")
total_written += written
logger.info(
"locale=%s: wrote %d new trajectories",
locale,
len(df),
written,
)

if total_written == 0:
Expand All @@ -591,13 +591,17 @@ def _simulate_to_parquet(
f"\nNext step: score them\n usersim eval --run {run_id} \\\n --trajectories {out} --out output/evaluations"
)

# Write the run manifest. ``run_id`` is the unix-epoch start time
# (see ``new_run_id``), so we pass it as ``started_at_epoch`` to keep
# the manifest's started/runtime fields meaningful even on resume.
# First-writer-wins: a resume into a run that already has a manifest
# is a noop. ``trajectories_requested`` is the per-locale request
# count from this invocation; the manifest's ``trajectories_completed``
# comes from re-reading the parquet partitions for this run.
if run_id == LEGACY_RUN_ID:
# The bare layout has no run folder to hold a manifest.
logger.info("run=%s is the bare legacy layout; no run manifest written", run_id)
return 0

# Write the run manifest. ``started_at_epoch`` is when the invocation
# that made the run started, kept in the run's record. First-writer-wins:
# a resume into a run that already has a manifest is a noop.
# ``trajectories_requested`` is the per-locale request count from this
# invocation; the manifest's ``trajectories_completed`` comes from
# re-reading the parquet partitions for this run.
requested_per_locale: dict[str, int] = {l: int(num_rows) for l in run_locales} if num_rows else {}
sim_cfg_for_manifest = _build_sim_config_for_manifest(
max_turns=max_turns,
Expand All @@ -607,7 +611,7 @@ def _simulate_to_parquet(
try:
manifest_path_written = write_simulator_manifest(
run_id=run_id,
started_at_epoch=int(run_id),
started_at_epoch=run_started_at(out, run_id) or int(started_at),
trajectory_root=out,
models_config=models,
models_config_path=models_path,
Expand Down
Loading
Loading