From f2989c5f25e3ac1804ec6f818ad30242f8ed10be Mon Sep 17 00:00:00 2001 From: isarasua Date: Fri, 9 Oct 2026 18:32:33 +0200 Subject: [PATCH 1/2] fix(storage): write simulation runs under a lock and publish them whole Two plain runs started in the same second shared one run folder, and two resumes of one run could each save a row the other had just saved. Every write into a trajectory root now holds one file lock per root (`filelock`, now declared, which excludes other threads as well as other processes). Where `flock` is unsupported, writing is refused before any model call, since `filelock`'s fallback locks can let two writers in. A new run appears only with its record (`_run_settings.json`) and first rows in place: both are written to a hidden folder that is then renamed, so a crash never leaves an empty run for eval or report to pick up, and nothing visible is ever deleted. Plain `simulate` resumes the latest run it made, and the manifest's start time comes from the run's record. Each batch's Data Designer artifacts go to a temporary folder of their own, removed once its rows are loaded, since Data Designer names them by the second a batch started. The simulate notebook saves through the same helpers. Co-Authored-By: Claude Opus 5.5 Signed-off-by: isarasua --- notebooks/01_simulate.ipynb | 138 +++++++----- pyproject.toml | 2 + src/usersim/cli/simulate.py | 97 ++++---- src/usersim/engine/core/storage.py | 270 +++++++++++++++++++++- tests/engine/core/test_run_writes.py | 274 +++++++++++++++++++++++ tests/test_cli_simulate_run_partition.py | 90 +++++++- tests/test_runs.py | 18 ++ uv.lock | 2 + 8 files changed, 778 insertions(+), 113 deletions(-) create mode 100644 tests/engine/core/test_run_writes.py diff --git a/notebooks/01_simulate.ipynb b/notebooks/01_simulate.ipynb index 776b740..5c4cc28 100644 --- a/notebooks/01_simulate.ipynb +++ b/notebooks/01_simulate.ipynb @@ -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", @@ -101,14 +103,17 @@ " preview_summary,\n", ")\n", "from usersim.engine.core.storage import (\n", + " LEGACY_RUN_ID,\n", + " check_run_write_lock,\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", @@ -577,12 +582,12 @@ "source": [ "## 🔁 Resumability check + run-id resolution\n", "\n", - "Every fresh simulation lands under its own `run=/` 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=/` 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=/`. Old runs are not deleted (they keep being readable via `read_partitioned_dataset(TRAJECTORY_PATH, run=\"\")` 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=/`, 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=\"\")` 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." ] @@ -593,6 +598,10 @@ "metadata": {}, "outputs": [], "source": [ + "# Refused here, before any model call, where TRAJECTORY_PATH's filesystem has\n", + "# no file locks: runs cannot be written there safely.\n", + "check_run_write_lock(TRAJECTORY_PATH)\n", + "\n", "if SIM_RUN_ID is not None:\n", " # Pinned run: finish a run you stopped part-way. RESUME reaches only the\n", " # LATEST run, so once a newer one exists the older partial is otherwise\n", @@ -619,20 +628,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", @@ -661,6 +671,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", @@ -684,59 +695,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}\")" ] }, { diff --git a/pyproject.toml b/pyproject.toml index bbe8327..6ee127d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,9 +57,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", diff --git a/src/usersim/cli/simulate.py b/src/usersim/cli/simulate.py index 71a144b..02059ff 100644 --- a/src/usersim/cli/simulate.py +++ b/src/usersim/cli/simulate.py @@ -27,9 +27,12 @@ import argparse import logging +import tempfile +import time from fractions import Fraction from pathlib import Path +from usersim.cli._errors import ConfigError from usersim.cli._pipeline import resolve_toolset_seed_for_mix, verify_probe_assets logger = logging.getLogger("usersim.cli.simulate") @@ -466,35 +469,35 @@ 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, + RunLockUnsupportedError, + check_run_write_lock, 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() + # Refused here, before any model call, rather than at the first save. + try: + check_run_write_lock(out) + except RunLockUnsupportedError as error: + raise ConfigError(str(error)) from None + + # ``--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", @@ -563,23 +566,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: @@ -591,13 +598,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, @@ -607,7 +618,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, diff --git a/src/usersim/engine/core/storage.py b/src/usersim/engine/core/storage.py index 64dcb88..5173e3b 100644 --- a/src/usersim/engine/core/storage.py +++ b/src/usersim/engine/core/storage.py @@ -11,10 +11,13 @@ 1. ``run=`` — top-level isolation per simulator invocation. Every fresh ``cli simulate`` (or notebook - ``RESUME=False``) mints a new run id; ``RESUME=True`` resumes - into the latest existing run. Eval and report stages inherit the + ``RESUME=False``) gets a new run id; ``RESUME=True`` resumes + into the latest run of its kind. Eval and report stages inherit the trajectory run's id so a single run is the unit of "X's - trajectories + X's evaluations + X's reports." + trajectories + X's evaluations + X's reports." A simulator writes + under ``run_write_lock``, and a new run appears only with its record + (``_run_settings.json``) and first rows in place; see "Run records + and locked writes" below. 2. ``locale=`` — one directory per locale, written atomically per-locale so a crash mid-run leaves earlier locales' partitions intact and resumable. @@ -44,12 +47,17 @@ from __future__ import annotations +import errno +import json import logging +import os import shutil import tempfile import time +import uuid +from contextlib import contextmanager from pathlib import Path -from typing import TYPE_CHECKING, Sequence +from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, Sequence if TYPE_CHECKING: # pyarrow is imported inside functions to keep it off the import path of @@ -188,6 +196,236 @@ def resolve_run_or_raise( return resolved +# --------------------------------------------------------------------------- +# Run records and locked writes +# +# A simulator writes into a trajectory root while holding ``run_write_lock``, +# and a new run becomes visible only with its record and its first rows in +# place, so two writers never share a run, a reader never sees an empty one, +# and the checks a writer makes against a run's saved rows still hold when it +# writes. +# --------------------------------------------------------------------------- + +#: What a run was made from, written inside ``run=/`` before the run is +#: visible. Data readers skip it, as pyarrow ignores files that start with ``_``. +RUN_RECORD_FILENAME: str = "_run_settings.json" + +#: The lock file every writer into a trajectory root holds. Dot-prefixed, so +#: neither ``list_runs`` nor pyarrow reads it. +RUN_LOCK_FILENAME: str = ".usersim.lock" + +#: 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" + +#: A new run is assembled in a hidden folder with this prefix, then renamed. +_CLAIM_PREFIX: str = ".claim-" + +#: Whether ``flock`` works under a trajectory root, found once per root. +_FLOCK_WORKS: dict[str, bool] = {} + + +class RunLockUnsupportedError(OSError): + """A trajectory root on a filesystem without ``flock``, where no run can be written safely.""" + + +def read_run_record(root: str | Path, run_id: str) -> dict[str, Any] | None: + """Return ``run_id``'s record, or ``None`` for a run written without one.""" + if run_id == LEGACY_RUN_ID: + return None + try: + return json.loads((run_subroot(root, run_id) / RUN_RECORD_FILENAME).read_text(encoding="utf-8")) + except FileNotFoundError: + return None + + +def latest_run_where( + root: str | Path, + keep: Callable[[str, dict[str, Any] | None], bool], +) -> str | None: + """Return the most recent run for which ``keep(run_id, record)`` is true. + + Runs are tried newest first, in ``latest_run_id``'s order, with + ``"legacy"`` last. + """ + runs = list_runs(root) + ordered = sorted((r for r in runs if r != LEGACY_RUN_ID), reverse=True) + ordered += [r for r in runs if r == LEGACY_RUN_ID] + for run_id in ordered: + if keep(run_id, read_run_record(root, run_id)): + return run_id + return None + + +def run_started_at(root: str | Path, run_id: str) -> int | None: + """When ``run_id`` started, in epoch seconds. + + The record's ``started_at`` (the start of the invocation that made the + run), else the run id itself for a numeric id from before records, else + ``None`` (the bare ``"legacy"`` layout). + """ + record = read_run_record(root, run_id) + if record is not None and record.get("started_at") is not None: + return int(record["started_at"]) + return int(run_id) if run_id.isdigit() else None + + +@contextmanager +def run_write_lock(root: str | Path) -> Iterator[Path]: + """Hold the write lock for trajectory root ``root``. + + Choosing a run, checking its saved rows and writing must all happen while + it is held, and nothing slow should: a simulator holds it per write, never + while it simulates. On entry it removes any hidden claim folder, which only + a writer that died while holding the lock can have left. + + It excludes other threads as well as other processes: each call takes its + own ``filelock.FileLock``, an ``flock`` lock on POSIX. It relies on locks + that every writing host shares, so on a mount whose ``flock`` locks are + local to one node (Lustre ``localflock``, NFS ``local_lock=flock``, + ``local_lock=all`` or ``nolock``), give each node its own root. Where + ``flock`` is not supported at all, it raises + :class:`RunLockUnsupportedError` rather than fall back to a lock file: + ``filelock``'s fallbacks can each let two writers in. + :func:`check_run_write_lock` finds that before any work starts. + """ + from filelock import FileLock + + root_path = check_run_write_lock(root) + with FileLock(root_path / RUN_LOCK_FILENAME): + for leftover in root_path.glob(f"{_CLAIM_PREFIX}*"): + shutil.rmtree(leftover, ignore_errors=True) + yield root_path + + +def check_run_write_lock(root: str | Path) -> Path: + """Create ``root`` if needed, and raise if runs cannot be written there safely. + + Raises :class:`RunLockUnsupportedError` when ``root``'s filesystem has no + ``flock``. Call it before work that ends in a write, such as simulating, + so a run is refused before any model call rather than at its first save. + Returns the resolved root. + """ + root_path = Path(root).resolve() + root_path.mkdir(parents=True, exist_ok=True) + if not _flock_works(root_path): + raise RunLockUnsupportedError( + f"file locks (flock) are not supported under {root_path}, so runs cannot be written there " + "safely; write them to a filesystem that supports file locks, such as a local disk or " + "Lustre mounted with the flock option" + ) + return root_path + + +def _flock_works(root: Path) -> bool: + """Whether ``flock`` works under ``root``, tried once on a scratch file there. + + Windows has no ``flock``; ``filelock`` locks with ``msvcrt`` there. + """ + try: + import fcntl + except ImportError: + return True + key = str(root) + if key not in _FLOCK_WORKS: + with tempfile.TemporaryFile(dir=root) as probe: + try: + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as error: + if error.errno != errno.ENOSYS: + raise + _FLOCK_WORKS[key] = False + else: + _FLOCK_WORKS[key] = True + return _FLOCK_WORKS[key] + + +def publish_new_run( + df, + root: str | Path, + *, + locale: str, + record: Mapping[str, Any], + clock: Callable[[], float] = time.time, + sleep: Callable[[float], None] = time.sleep, +) -> str: + """Make a new run visible with its record and its first rows, and return its id. + + Call while holding :func:`run_write_lock`. The record and ``df``'s + ``locale`` rows are written into a hidden folder that is then renamed to + ``run=``, the current second. If that id is taken, it waits for the + next second rather than take one a later run would mint for itself. + Renaming onto a run folder that holds anything fails, and every run made + here holds its record from the start, so two runs never share a folder. + """ + root_path = Path(root).resolve() + # Made like any folder, so the run gets the usual permissions: the rename + # keeps them, and a ``tempfile.mkdtemp`` folder would be owner-only. + claim = root_path / f"{_CLAIM_PREFIX}{uuid.uuid4().hex}" + claim.mkdir() + published = False + try: + (claim / RUN_RECORD_FILENAME).write_text(json.dumps(dict(record), sort_keys=True), encoding="utf-8") + if _write_locale_rows(df, claim, locale=locale) is None: + raise ValueError(f"publish_new_run: no {locale} rows to publish; a run is never made empty") + while True: + run_id = new_run_id(clock()) + try: + os.rename(claim, run_subroot(root_path, run_id)) + except OSError as error: + if error.errno not in (errno.EEXIST, errno.ENOTEMPTY): + raise + sleep(max(0.0, int(run_id) + 1 - clock())) + continue + published = True + return run_id + finally: + if not published: + shutil.rmtree(claim, ignore_errors=True) + + +def sampled_run_to_resume(root: str | Path) -> str | None: + """The run a resumed simulation of sampled rows reopens: the latest one of its kind. + + A run made before runs recorded their kind counts as one; a run of any + other kind never does. + """ + return latest_run_where(root, lambda run_id, record: record is None or record.get("kind") == RUN_KIND_SAMPLED) + + +def save_sampled_rows( + df, + root: str | Path, + *, + locale: str, + run_id: str | None, + started_at: float, +) -> tuple[str | None, int]: + """Save one locale's sampled rows into ``run_id``, or into a new run when it is ``None``. + + Under the write lock, rows already saved in ``run_id`` are skipped, + whoever saved them, so no id is saved twice in a run; a new run never + shares a folder with another; and a run of another kind is refused, so a + run holds one kind of rows. Returns the run the rows went to and how many + were written. + """ + with run_write_lock(root): + if run_id is not None: + record = read_run_record(root, run_id) + if record is not None and record.get("kind") != RUN_KIND_SAMPLED: + raise ValueError(f"run {run_id} holds {record.get('kind')} rows; sampled rows need a run of their own") + df = df[~df["trajectory_id"].isin(existing_trajectory_ids(root, run=run_id))] + df = df[df["locale"].astype(str) == str(locale)] + if df.empty: + return run_id, 0 + if run_id is None: + record = {"kind": RUN_KIND_SAMPLED, "started_at": started_at} + run_id = publish_new_run(df, root, locale=locale, record=record) + else: + write_locale_partition(df, root, locale=locale, run_id=run_id) + return run_id, len(df) + + def write_run_partition( df, root: str | Path, @@ -996,20 +1234,36 @@ def write_locale_partition( batch leaves the prior locales' partitions on disk and resumable (within the same run id). """ + written = _write_locale_rows( + df, + run_subroot(root, run_id), + locale=locale, + partition_cols=partition_cols, + writer_id=writer_id, + ) + return written if written is not None else Path(root).resolve() + + +def _write_locale_rows( + df, + subroot: Path, + *, + locale: str, + partition_cols: Sequence[str] = DEFAULT_PARTITION_COLS, + writer_id: str | None = None, +) -> Path | None: + """Write ``df``'s ``locale`` rows under ``subroot``; ``None`` when there are none.""" if "locale" not in df.columns: raise ValueError("write_locale_partition: dataframe is missing 'locale' column") locale_mask = df["locale"].astype(str) == str(locale) locale_df = df[locale_mask] if locale_df.empty: - return Path(root).resolve() - - import os + return None if writer_id is None: writer_id = f"pid{os.getpid()}" basename_template = f"part-{{i}}-{locale}-{writer_id}.parquet" - subroot = run_subroot(root, run_id) return write_partitioned_dataset( locale_df, subroot, diff --git a/tests/engine/core/test_run_writes.py b/tests/engine/core/test_run_writes.py new file mode 100644 index 0000000..f086ea6 --- /dev/null +++ b/tests/engine/core/test_run_writes.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runs are written under one lock per trajectory root, and appear whole. + +The invariants: a run becomes visible only with its record and its first rows +in place, two writers never share a run, nothing visible is ever deleted, and +plain ``simulate`` never writes into a run of another kind. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import textwrap +import threading +import time +from pathlib import Path + +import pandas as pd +import pytest + +from usersim.engine.core.storage import ( + LEGACY_RUN_ID, + RUN_KIND_SAMPLED, + RUN_RECORD_FILENAME, + RunLockUnsupportedError, + check_run_write_lock, + existing_trajectory_ids, + latest_run_where, + list_runs, + publish_new_run, + read_partitioned_dataset, + read_run_record, + run_started_at, + run_write_lock, + sampled_run_to_resume, + save_sampled_rows, +) + +#: A kind other than sampled rows, as a later kind of run would record. +_OTHER_KIND = "other" + + +def _rows(*ids: str, locale: str = "en_US") -> pd.DataFrame: + return pd.DataFrame( + {"trajectory_id": list(ids), "locale": [locale] * len(ids), "probe_family": ["general"] * len(ids)} + ) + + +def _publish(root: Path, *ids: str, kind: str = RUN_KIND_SAMPLED, now: float = 1_800_000_000.0) -> str: + with run_write_lock(root): + return publish_new_run(_rows(*ids), root, locale="en_US", record={"kind": kind}, clock=lambda: now) + + +def test_a_new_run_appears_with_its_record_and_rows(tmp_path): + run_id = _publish(tmp_path, "a", "b") + + assert list_runs(tmp_path) == [run_id] + assert read_run_record(tmp_path, run_id) == {"kind": RUN_KIND_SAMPLED} + assert existing_trajectory_ids(tmp_path, run=run_id) == {"a", "b"} + assert not list(tmp_path.glob(".claim-*")) + + +def test_a_taken_second_waits_for_the_next_one(tmp_path): + taken = _publish(tmp_path, "a", now=1_800_000_000.0) + ticks = iter([1_800_000_000.2, 1_800_000_000.2, 1_800_000_001.0]) + waits: list[float] = [] + + with run_write_lock(tmp_path): + run_id = publish_new_run( + _rows("b"), + tmp_path, + locale="en_US", + record={"kind": RUN_KIND_SAMPLED}, + clock=lambda: next(ticks), + sleep=waits.append, + ) + + assert (taken, run_id) == ("1800000000", "1800000001") + assert waits == [pytest.approx(0.8)] + assert existing_trajectory_ids(tmp_path, run=taken) == {"a"} + + +def test_a_run_is_never_published_empty(tmp_path): + with run_write_lock(tmp_path), pytest.raises(ValueError, match="never made empty"): + publish_new_run(_rows(), tmp_path, locale="en_US", record={"kind": RUN_KIND_SAMPLED}) + + assert list_runs(tmp_path) == [] + assert not list(tmp_path.glob(".claim-*")) + + +def test_a_failed_write_leaves_nothing_visible(tmp_path): + with run_write_lock(tmp_path), pytest.raises(ValueError): + publish_new_run(pd.DataFrame({"trajectory_id": ["a"]}), tmp_path, locale="en_US", record={}) + + assert list_runs(tmp_path) == [] + assert not list(tmp_path.glob(".claim-*")) + + +def test_the_lock_clears_a_claim_a_dead_writer_left(tmp_path): + abandoned = tmp_path / ".claim-abandoned" + abandoned.mkdir() + (abandoned / RUN_RECORD_FILENAME).write_text("{}") + + with run_write_lock(tmp_path): + pass + + assert not abandoned.exists() + + +def test_the_lock_excludes_another_process(tmp_path): + """A second process waits until the first releases the lock.""" + holder = subprocess.Popen( + [ + sys.executable, + "-c", + textwrap.dedent( + f""" + import time + from usersim.engine.core.storage import run_write_lock + with run_write_lock({str(tmp_path)!r}): + print("locked", flush=True) + time.sleep(1.0) + """ + ), + ], + stdout=subprocess.PIPE, + text=True, + ) + try: + assert holder.stdout.readline().strip() == "locked" + started = time.monotonic() + with run_write_lock(tmp_path): + waited = time.monotonic() - started + finally: + holder.wait(timeout=30) + + assert waited > 0.5 + + +def test_runs_are_searched_newest_first_and_legacy_last(tmp_path): + (tmp_path / "locale=en_US").mkdir() + older = _publish(tmp_path, "a", kind=RUN_KIND_SAMPLED, now=1_700_000_000.0) + newer = _publish(tmp_path, "b", kind=_OTHER_KIND, now=1_800_000_000.0) + + assert latest_run_where(tmp_path, lambda run, record: True) == newer + assert latest_run_where(tmp_path, lambda run, record: record and record["kind"] == RUN_KIND_SAMPLED) == older + + +def test_a_run_started_when_its_invocation_did(tmp_path): + with run_write_lock(tmp_path): + run_id = publish_new_run( + _rows("a"), + tmp_path, + locale="en_US", + record={"kind": RUN_KIND_SAMPLED, "started_at": 1_700_000_000.4}, + clock=lambda: 1_700_000_100.0, + ) + (tmp_path / "run=1600000000").mkdir() + + assert run_id == "1700000100" + assert run_started_at(tmp_path, run_id) == 1_700_000_000 + assert run_started_at(tmp_path, "1600000000") == 1_600_000_000 + assert run_started_at(tmp_path, LEGACY_RUN_ID) is None + + +# --------------------------------------------------------------------------- +# Plain ``simulate`` +# --------------------------------------------------------------------------- + + +def test_plain_resume_never_reopens_a_run_of_another_kind(tmp_path): + sampled = _publish(tmp_path, "a", kind=RUN_KIND_SAMPLED, now=1_700_000_000.0) + _publish(tmp_path, "b", kind=_OTHER_KIND, now=1_800_000_000.0) + + assert sampled_run_to_resume(tmp_path) == sampled + + +def test_plain_resume_reopens_a_run_made_before_runs_had_records(tmp_path): + _publish(tmp_path, "a", kind=_OTHER_KIND, now=1_800_000_000.0) + old = tmp_path / "run=1700000000" + old.mkdir() + + assert sampled_run_to_resume(tmp_path) == "1700000000" + + +def test_a_fresh_plain_run_never_joins_a_run_started_in_the_same_second(tmp_path): + other = _publish(tmp_path, "s", kind=_OTHER_KIND, now=time.time()) + + run_id, written = save_sampled_rows(_rows("p"), tmp_path, locale="en_US", run_id=None, started_at=time.time()) + + assert written == 1 and run_id != other + assert existing_trajectory_ids(tmp_path, run=other) == {"s"} + assert read_run_record(tmp_path, run_id)["kind"] == RUN_KIND_SAMPLED + + +def test_rows_saved_in_the_run_since_it_was_chosen_are_not_saved_twice(tmp_path): + """Whoever saved them: a resume, or the writer that made the run, as with ``--no-skip-existing``.""" + run_id = _publish(tmp_path, "a", kind=RUN_KIND_SAMPLED) + + _, written = save_sampled_rows(_rows("a", "b"), tmp_path, locale="en_US", run_id=run_id, started_at=0.0) + + assert written == 1 + assert sorted(read_partitioned_dataset(tmp_path, run=run_id)["trajectory_id"]) == ["a", "b"] + + +def test_a_fresh_plain_run_records_when_it_started(tmp_path): + run_id, _ = save_sampled_rows(_rows("a"), tmp_path, locale="en_US", run_id=None, started_at=1_700_000_000.0) + + assert json.loads((tmp_path / f"run={run_id}" / RUN_RECORD_FILENAME).read_text()) == { + "kind": RUN_KIND_SAMPLED, + "started_at": 1_700_000_000.0, + } + + +def test_sampled_rows_never_go_into_a_run_of_another_kind(tmp_path): + other = _publish(tmp_path, "s", kind=_OTHER_KIND) + + with pytest.raises(ValueError, match="run of their own"): + save_sampled_rows(_rows("p"), tmp_path, locale="en_US", run_id=other, started_at=0.0) + assert existing_trajectory_ids(tmp_path, run=other) == {"s"} + + +def test_a_new_run_gets_the_usual_permissions(tmp_path): + """Others sharing the root can read a run, as they could before runs were published.""" + previous = os.umask(0o022) + try: + run_id = _publish(tmp_path, "a") + finally: + os.umask(previous) + + assert (tmp_path / f"run={run_id}").stat().st_mode & 0o777 == 0o755 + + +def test_the_lock_excludes_another_thread(tmp_path): + """Two threads of one process exclude each other, as two processes do.""" + held = threading.Event() + + def hold(): + with run_write_lock(tmp_path): + held.set() + time.sleep(1.0) + + holder = threading.Thread(target=hold) + holder.start() + try: + assert held.wait(timeout=30) + started = time.monotonic() + with run_write_lock(tmp_path): + waited = time.monotonic() - started + finally: + holder.join(timeout=30) + + assert waited > 0.5 + + +def test_where_flock_is_unsupported_writing_is_refused(tmp_path, monkeypatch): + """No fallback lock: ``filelock``'s can each let two writers in, so a root without ``flock`` takes no runs.""" + 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 pytest.raises(RunLockUnsupportedError, match="not supported"): + check_run_write_lock(tmp_path) + with pytest.raises(RunLockUnsupportedError), run_write_lock(tmp_path): + pass + assert list(tmp_path.iterdir()) == [] diff --git a/tests/test_cli_simulate_run_partition.py b/tests/test_cli_simulate_run_partition.py index 8da7ab9..0bea142 100644 --- a/tests/test_cli_simulate_run_partition.py +++ b/tests/test_cli_simulate_run_partition.py @@ -38,6 +38,7 @@ from unittest.mock import patch import pandas as pd +import pytest from usersim.engine.core._assets import packaged_assets_dir from usersim.engine.core.storage import ( @@ -244,7 +245,7 @@ def load_dataset(self): return _synthetic_trajectory("FRESH") class _FakeDD: - def create(self, builder, num_records=None): + def create(self, builder, num_records=None, artifact_path=None): return _FakeResult() captured_builder_kwargs = {} @@ -332,7 +333,7 @@ def load_dataset(self): return _synthetic_trajectory("RESUMED") class _FakeDD: - def create(self, builder, num_records=None): + def create(self, builder, num_records=None, artifact_path=None): return _FakeResult() def _fake_builder(**kwargs): @@ -376,7 +377,7 @@ def load_dataset(self): return _synthetic_context_failure() class _FakeDD: - def create(self, builder, num_records=None): + def create(self, builder, num_records=None, artifact_path=None): assert num_records == 1 return _FakeResult() @@ -407,3 +408,86 @@ def create(self, builder, num_records=None): assert row["finance_task_id"] == "TASK-CTX" assert json.loads(row["retrieved_document_ids"]) == ["d1"] assert "evidence" in row["conversation_messages"] + + +def test_each_batch_gets_its_own_artifacts_and_the_manifest_its_start(tmp_path: Path, monkeypatch) -> None: + """Data Designer names a batch's artifacts by the second it started, so two runs in one directory + could share them. Each batch gets a folder of its own, removed once its rows are loaded, and the + manifest's start is when the invocation started, kept in the run's record. + """ + from types import SimpleNamespace + + from usersim.cli import simulate + from usersim.cli._models import default_models_path, load_models_config + from usersim.engine.core.manifest import load_run_manifest + + folders: list[Path] = [] + + class _FakeResult: + def load_dataset(self): + return _synthetic_trajectory("FRESH") + + class _FakeDD: + def create(self, builder, num_records=None, artifact_path=None): + folders.append(Path(artifact_path)) + assert folders[-1].is_dir() + return _FakeResult() + + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(simulate, "time", SimpleNamespace(time=lambda: 1_700_000_000.4)) + out = tmp_path / "trajs" + with patch("usersim.cli._pipeline.build_simulator_config_builder", lambda **kwargs: (_FakeDD(), object())): + rc = simulate._simulate_to_parquet( + models=load_models_config(default_models_path()), + models_path=default_models_path(), + panel=None, + locales=["en_US"], + num_rows=1, + out=out, + assets_dir=packaged_assets_dir(), + max_turns=1, + probe_mix={"general_open_ended": 1.0}, + random_seed=42, + skip_existing=True, + ) + + assert rc == 0 + (run_id,) = list_runs(out) + assert load_run_manifest(run_id, out).started_at_epoch == 1_700_000_000 + assert len(folders) == 1 and not folders[0].exists() + assert tmp_path not in folders[0].parents + + +def test_a_root_without_file_locks_is_refused_before_any_model_call(tmp_path: Path, monkeypatch) -> None: + """The user is told which folder and why, as a configuration error, before anything is simulated.""" + fcntl = pytest.importorskip("fcntl") + import errno + + from usersim.cli import simulate + from usersim.cli._errors import ConfigError + from usersim.cli._models import default_models_path, load_models_config + + def unsupported(*args, **kwargs): + raise OSError(errno.ENOSYS, "flock is not supported here") + + def no_pipeline(**kwargs): + raise AssertionError("a pipeline was built") + + monkeypatch.setattr(fcntl, "flock", unsupported) + with ( + patch("usersim.cli._pipeline.build_simulator_config_builder", no_pipeline), + pytest.raises(ConfigError, match="file locks"), + ): + simulate._simulate_to_parquet( + models=load_models_config(default_models_path()), + models_path=default_models_path(), + panel=None, + locales=["en_US"], + num_rows=1, + out=tmp_path / "trajs", + assets_dir=packaged_assets_dir(), + max_turns=1, + probe_mix={"general_open_ended": 1.0}, + random_seed=42, + skip_existing=True, + ) diff --git a/tests/test_runs.py b/tests/test_runs.py index 7d54230..e5e0a7e 100644 --- a/tests/test_runs.py +++ b/tests/test_runs.py @@ -296,6 +296,24 @@ def test_simulate_notebook_warns_about_already_finished_locales() -> None: assert "re-simulated at full LLM cost" in src +def test_simulate_notebook_writes_runs_as_the_cli_does() -> None: + """The notebook saves through the same locked helpers as ``usersim simulate``. + + Writing partitions itself would let it resume a run of another kind, share + a run with another writer started in the same second, or save an id twice. + """ + src = _simulate_notebook_source() + # Refused before any model call where the trajectory root has no file locks. + assert "check_run_write_lock(TRAJECTORY_PATH)" in src + assert "sampled_run_to_resume(TRAJECTORY_PATH)" in src + assert "save_sampled_rows(" in src + assert "write_locale_partition(" not in src and "new_run_id(" not in src + # Each batch's Data Designer artifacts in a folder of its own. + assert "artifact_path=_artifacts" in src + # The bare legacy layout has no run folder to hold a manifest. + assert "elif RUN_ID == LEGACY_RUN_ID:" in src + + def test_resumed_manifest_describes_the_run_not_the_last_pass() -> None: """The manifest is first-writer-wins, so the pass that happens to complete a run writes the only manifest it will ever have. Recording `locales_requested=LOCALES` diff --git a/uv.lock b/uv.lock index 4dea804..3a047a2 100644 --- a/uv.lock +++ b/uv.lock @@ -1947,6 +1947,7 @@ version = "0.1.0" source = { editable = "." } dependencies = [ { name = "data-designer" }, + { name = "filelock" }, { name = "httpx" }, { name = "huggingface-hub" }, { name = "langdetect" }, @@ -1978,6 +1979,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "data-designer", specifier = ">=0.9.4,<0.10" }, + { name = "filelock" }, { name = "httpx" }, { name = "huggingface-hub" }, { name = "langdetect" }, From f4045c974dd4c5b07d8b49b996d814b458064de4 Mon Sep 17 00:00:00 2001 From: isarasua Date: Mon, 12 Oct 2026 21:51:32 +0200 Subject: [PATCH 2/2] fix(storage): write where flock is unsupported, and keep live writers' claims Addresses the review: - Where `flock` is unsupported (ENOSYS on Lustre mounted without it, ENOLCK on NFS without its lock service, EOPNOTSUPP on some network mounts), runs are written without the cross-process lock instead of being refused, and a warning naming what that loses is logged once per root. The threads of one process still take turns, `filelock`'s soft-lock fallback is not used, and the up-front checks in `simulate` and the notebook are gone. - Taking the lock removes a claim folder only when its writer died: it was made in this process's PID scope (the boot and PID namespace, on Linux), its process is gone, and it is over an hour old. Any other claim is kept, so a slow, paused or remote writer never loses its batch. - Every part file is written under a hidden name and renamed once complete, so another writer's scan never opens a half-written one. Co-Authored-By: Claude Opus 5.5 Signed-off-by: isarasua --- notebooks/01_simulate.ipynb | 5 - src/usersim/cli/simulate.py | 9 +- src/usersim/engine/core/storage.py | 184 ++++++++++++++++++----- tests/engine/core/test_run_writes.py | 161 ++++++++++++++++++-- tests/test_cli_simulate_run_partition.py | 30 ++-- tests/test_runs.py | 2 - 6 files changed, 313 insertions(+), 78 deletions(-) diff --git a/notebooks/01_simulate.ipynb b/notebooks/01_simulate.ipynb index 5c4cc28..3071a01 100644 --- a/notebooks/01_simulate.ipynb +++ b/notebooks/01_simulate.ipynb @@ -104,7 +104,6 @@ ")\n", "from usersim.engine.core.storage import (\n", " LEGACY_RUN_ID,\n", - " check_run_write_lock,\n", " existing_trajectory_ids,\n", " latest_run_id,\n", " list_partition_keys,\n", @@ -598,10 +597,6 @@ "metadata": {}, "outputs": [], "source": [ - "# Refused here, before any model call, where TRAJECTORY_PATH's filesystem has\n", - "# no file locks: runs cannot be written there safely.\n", - "check_run_write_lock(TRAJECTORY_PATH)\n", - "\n", "if SIM_RUN_ID is not None:\n", " # Pinned run: finish a run you stopped part-way. RESUME reaches only the\n", " # LATEST run, so once a newer one exists the older partial is otherwise\n", diff --git a/src/usersim/cli/simulate.py b/src/usersim/cli/simulate.py index 02059ff..dc520e7 100644 --- a/src/usersim/cli/simulate.py +++ b/src/usersim/cli/simulate.py @@ -32,7 +32,6 @@ from fractions import Fraction from pathlib import Path -from usersim.cli._errors import ConfigError from usersim.cli._pipeline import resolve_toolset_seed_for_mix, verify_probe_assets logger = logging.getLogger("usersim.cli.simulate") @@ -470,8 +469,6 @@ def _simulate_to_parquet( from usersim.engine.core.manifest import write_simulator_manifest from usersim.engine.core.storage import ( LEGACY_RUN_ID, - RunLockUnsupportedError, - check_run_write_lock, existing_trajectory_ids, run_started_at, sampled_run_to_resume, @@ -480,11 +477,7 @@ def _simulate_to_parquet( started_at = time.time() out = Path(out).resolve() - # Refused here, before any model call, rather than at the first save. - try: - check_run_write_lock(out) - except RunLockUnsupportedError as error: - raise ConfigError(str(error)) from None + out.mkdir(parents=True, exist_ok=True) # ``--no-skip-existing`` starts a fresh run; the default resumes the # latest run this command made. A fresh run is claimed only when its diff --git a/src/usersim/engine/core/storage.py b/src/usersim/engine/core/storage.py index 5173e3b..536820c 100644 --- a/src/usersim/engine/core/storage.py +++ b/src/usersim/engine/core/storage.py @@ -51,8 +51,10 @@ import json import logging import os +import re import shutil import tempfile +import threading import time import uuid from contextlib import contextmanager @@ -201,9 +203,9 @@ def resolve_run_or_raise( # # A simulator writes into a trajectory root while holding ``run_write_lock``, # and a new run becomes visible only with its record and its first rows in -# place, so two writers never share a run, a reader never sees an empty one, -# and the checks a writer makes against a run's saved rows still hold when it -# writes. +# place, so two writers never share a run and a reader never sees an empty +# one. Where the lock is shared by every writer, the checks a writer makes +# against a run's saved rows also still hold when it writes. # --------------------------------------------------------------------------- #: What a run was made from, written inside ``run=/`` before the run is @@ -221,12 +223,30 @@ def resolve_run_or_raise( #: A new run is assembled in a hidden folder with this prefix, then renamed. _CLAIM_PREFIX: str = ".claim-" +#: A claim folder is removed only once it is older than this, and its writer +#: is gone; a live one lasts only while one locale's rows are written and +#: renamed. +_STALE_CLAIM_SECONDS: float = 3600.0 + #: Whether ``flock`` works under a trajectory root, found once per root. _FLOCK_WORKS: dict[str, bool] = {} +#: What ``flock`` fails with where a filesystem has no usable locks: not +#: implemented (Lustre without ``flock``), no lock service (NFS without +#: ``lockd``), or not supported (some network and FUSE mounts). +_NO_FLOCK_ERRNOS = frozenset({errno.ENOSYS, errno.ENOLCK, errno.EOPNOTSUPP, errno.ENOTSUP}) + +#: Where Linux names this machine's current boot, and this process's PID +#: namespace: together, the scope in which a PID identifies one process. +_BOOT_ID_PATH = Path("/proc/sys/kernel/random/boot_id") +_PID_NAMESPACE_PATH = Path("/proc/self/ns/pid") + +#: Roots already warned about writing without ``flock``. +_WARNED_WITHOUT_FLOCK: set[str] = set() -class RunLockUnsupportedError(OSError): - """A trajectory root on a filesystem without ``flock``, where no run can be written safely.""" +#: Where ``flock`` is unsupported, the writers of this process still take +#: turns through this lock. +_IN_PROCESS_WRITES = threading.Lock() def read_run_record(root: str | Path, run_id: str) -> dict[str, Any] | None: @@ -276,45 +296,121 @@ def run_write_lock(root: str | Path) -> Iterator[Path]: Choosing a run, checking its saved rows and writing must all happen while it is held, and nothing slow should: a simulator holds it per write, never - while it simulates. On entry it removes any hidden claim folder, which only - a writer that died while holding the lock can have left. + while it simulates. On entry it removes the claim folders whose writer + died: claims made in this process's PID scope (this boot and PID + namespace, on Linux), whose process is gone, and that are over an hour + old. A live or paused writer's claim, and any claim made in another + scope or where the scope can't be read, are left alone; readers never + see a claim either way. It excludes other threads as well as other processes: each call takes its - own ``filelock.FileLock``, an ``flock`` lock on POSIX. It relies on locks - that every writing host shares, so on a mount whose ``flock`` locks are - local to one node (Lustre ``localflock``, NFS ``local_lock=flock``, - ``local_lock=all`` or ``nolock``), give each node its own root. Where - ``flock`` is not supported at all, it raises - :class:`RunLockUnsupportedError` rather than fall back to a lock file: - ``filelock``'s fallbacks can each let two writers in. - :func:`check_run_write_lock` finds that before any work starts. + own ``filelock.FileLock``, an ``flock`` lock on POSIX. Where ``flock`` is + not supported (as on Lustre mounted without ``flock``, or NFS without its + lock service), only the threads of this process take turns, and a warning + is logged once per root. ``filelock``'s own fallback to a soft lock is not used: it can let + two writers in. On a mount whose ``flock`` locks are local to one node + (Lustre ``localflock``, NFS ``local_lock=flock``, ``local_lock=all`` or + ``nolock``), writers on different nodes are not excluded either. + + 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. """ from filelock import FileLock - root_path = check_run_write_lock(root) - with FileLock(root_path / RUN_LOCK_FILENAME): - for leftover in root_path.glob(f"{_CLAIM_PREFIX}*"): - shutil.rmtree(leftover, ignore_errors=True) + root_path = Path(root).resolve() + root_path.mkdir(parents=True, exist_ok=True) + shared = _flock_works(root_path) + with FileLock(root_path / RUN_LOCK_FILENAME) if shared else _IN_PROCESS_WRITES: + if not shared: + _warn_without_flock(root_path) + _remove_stale_claims(root_path) yield root_path -def check_run_write_lock(root: str | Path) -> Path: - """Create ``root`` if needed, and raise if runs cannot be written there safely. +def _pid_scope() -> str | None: + """The scope in which this process's PIDs identify processes: this boot and this PID namespace. - Raises :class:`RunLockUnsupportedError` when ``root``'s filesystem has no - ``flock``. Call it before work that ends in a write, such as simulating, - so a run is refused before any model call rather than at its first save. - Returns the resolved root. + One process can tell whether another is alive only within one scope, so + a claim made in another (another host, a container with its own PID + namespace, or before a reboot) is never judged from here. ``None`` where + the scope cannot be read: outside Linux, or with ``/proc`` hidden. """ - root_path = Path(root).resolve() - root_path.mkdir(parents=True, exist_ok=True) - if not _flock_works(root_path): - raise RunLockUnsupportedError( - f"file locks (flock) are not supported under {root_path}, so runs cannot be written there " - "safely; write them to a filesystem that supports file locks, such as a local disk or " - "Lustre mounted with the flock option" - ) - return root_path + try: + boot = _BOOT_ID_PATH.read_text(encoding="ascii").strip().replace("-", "") + namespace = os.stat(_PID_NAMESPACE_PATH).st_ino + except (OSError, ValueError): + return None + if not re.fullmatch(r"[0-9a-f]{32}", boot): + return None + return f"{boot}.{namespace}" + + +def _claim_name() -> str: + """A new claim folder's name, recording its writer's PID scope and PID where the scope is known.""" + scope = _pid_scope() + owner = f"{scope}-{os.getpid()}-" if scope else "" + return f"{_CLAIM_PREFIX}{owner}{uuid.uuid4().hex}" + + +def _claim_owner(name: str) -> tuple[str, int] | None: + """The PID scope and PID a claim folder's name records, or ``None`` if it records none.""" + match = re.fullmatch(r"([0-9a-f]{32}\.\d+)-(\d+)-[0-9a-f]{32}", name[len(_CLAIM_PREFIX) :]) + return (match.group(1), int(match.group(2))) if match else None + + +def _process_alive(pid: int) -> bool: + """Whether process ``pid`` exists in this PID scope; assumed alive where that cannot be checked.""" + if os.name == "nt": + # ``os.kill`` on Windows signals or terminates rather than probes. + return True + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except OSError: + return True + return True + + +def _remove_stale_claims(root: Path) -> None: + """Remove the claim folders under ``root`` whose writer, in this PID scope, is gone. + + A claim is removed only when it was made in this process's PID scope + (:func:`_pid_scope`), its process no longer exists, and it is older than + ``_STALE_CLAIM_SECONDS``. Any other claim is kept, since whether its + writer lives cannot be told from here. A dead writer can no longer rename + its claim, so removing one never races a publication. + """ + here = _pid_scope() + if here is None: + return + cutoff = time.time() - _STALE_CLAIM_SECONDS + for claim in root.glob(f"{_CLAIM_PREFIX}*"): + owner = _claim_owner(claim.name) + if owner is None or owner[0] != here or _process_alive(owner[1]): + continue + try: + stale = claim.stat().st_mtime < cutoff + except FileNotFoundError: + continue + if stale: + shutil.rmtree(claim, ignore_errors=True) + + +def _warn_without_flock(root: Path) -> None: + """Warn, once per root, that runs under ``root`` are written without a shared lock. + + Called while holding the in-process lock, so threads never warn twice. + """ + if str(root) in _WARNED_WITHOUT_FLOCK: + return + _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.", + root, + ) def _flock_works(root: Path) -> bool: @@ -332,7 +428,7 @@ def _flock_works(root: Path) -> bool: try: fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) except OSError as error: - if error.errno != errno.ENOSYS: + if error.errno not in _NO_FLOCK_ERRNOS: raise _FLOCK_WORKS[key] = False else: @@ -361,7 +457,7 @@ def publish_new_run( root_path = Path(root).resolve() # Made like any folder, so the run gets the usual permissions: the rename # keeps them, and a ``tempfile.mkdtemp`` folder would be owner-only. - claim = root_path / f"{_CLAIM_PREFIX}{uuid.uuid4().hex}" + claim = root_path / _claim_name() claim.mkdir() published = False try: @@ -404,7 +500,8 @@ def save_sampled_rows( """Save one locale's sampled rows into ``run_id``, or into a new run when it is ``None``. Under the write lock, rows already saved in ``run_id`` are skipped, - whoever saved them, so no id is saved twice in a run; a new run never + whoever saved them, so where the lock is shared no id is saved twice in a + run (see :func:`run_write_lock` for where it is not); a new run never shares a folder with another; and a run of another kind is refused, so a run holds one kind of rows. Returns the run the rows went to and how many were written. @@ -639,10 +736,19 @@ def write_partitioned_dataset( pass elif existing_data_behavior == "error" and out_path.exists(): raise FileExistsError(out_path) - # Subset to non-partition columns + write. + # Subset to non-partition columns + write. Each file is written under + # a hidden name and renamed once complete, so a reader (another + # writer's scan included) never opens a partial one. The name is not + # from ``mkstemp``, which would make the file readable by its owner + # only. sub_df = group_df[non_partition_cols] table = _arrow_table_for_write(sub_df) - pq.write_table(table, out_path) + partial = out_dir / f".{filename}.{uuid.uuid4().hex}.partial" + try: + pq.write_table(table, partial) + os.replace(partial, out_path) + finally: + partial.unlink(missing_ok=True) return root_path diff --git a/tests/engine/core/test_run_writes.py b/tests/engine/core/test_run_writes.py index f086ea6..ed7b410 100644 --- a/tests/engine/core/test_run_writes.py +++ b/tests/engine/core/test_run_writes.py @@ -10,24 +10,27 @@ from __future__ import annotations +import errno import json +import logging import os +import re import subprocess import sys import textwrap import threading import time +import uuid from pathlib import Path import pandas as pd import pytest +from usersim.engine.core import storage from usersim.engine.core.storage import ( LEGACY_RUN_ID, RUN_KIND_SAMPLED, RUN_RECORD_FILENAME, - RunLockUnsupportedError, - check_run_write_lock, existing_trajectory_ids, latest_run_where, list_runs, @@ -38,6 +41,7 @@ run_write_lock, sampled_run_to_resume, save_sampled_rows, + write_locale_partition, ) #: A kind other than sampled rows, as a later kind of run would record. @@ -100,15 +104,105 @@ def test_a_failed_write_leaves_nothing_visible(tmp_path): assert not list(tmp_path.glob(".claim-*")) -def test_the_lock_clears_a_claim_a_dead_writer_left(tmp_path): - abandoned = tmp_path / ".claim-abandoned" - abandoned.mkdir() - (abandoned / RUN_RECORD_FILENAME).write_text("{}") +#: Two PID scopes on one boot: this process's, and a container's with its own PID namespace. +_SCOPE = "0123456789abcdef0123456789abcdef.4026531836" +_OTHER_NAMESPACE = "0123456789abcdef0123456789abcdef.4026532001" + + +def _claim(root: Path, *, scope: str, pid: int, age_seconds: float) -> Path: + """A claim folder made by ``pid`` in ``scope``, last modified ``age_seconds`` ago.""" + claim = root / f".claim-{scope}-{pid}-{uuid.uuid4().hex}" + claim.mkdir() + (claim / RUN_RECORD_FILENAME).write_text("{}") + then = time.time() - age_seconds + os.utime(claim, (then, then)) + return claim + + +def _dead_pid() -> int: + finished = subprocess.Popen([sys.executable, "-c", "pass"]) + finished.wait() + return finished.pid + + +def test_the_pid_scope_is_this_boot_and_pid_namespace_where_linux_shows_them(): + scope = storage._pid_scope() + if Path("/proc/self/ns/pid").exists() and Path("/proc/sys/kernel/random/boot_id").exists(): + assert re.fullmatch(r"[0-9a-f]{32}\.\d+", scope or "") + else: + assert scope is None + + +@pytest.mark.skipif(os.name == "nt", reason="a PID is only probed where os.kill can probe one") +def test_the_lock_clears_only_a_claim_whose_writer_died_in_this_pid_scope(tmp_path, monkeypatch): + """A slow or paused writer's claim survives, and so does one whose PID this process cannot see.""" + monkeypatch.setattr(storage, "_pid_scope", lambda: _SCOPE) + two_hours = 2 * 3600 + abandoned = _claim(tmp_path, scope=_SCOPE, pid=_dead_pid(), age_seconds=two_hours) + just_left = _claim(tmp_path, scope=_SCOPE, pid=_dead_pid(), age_seconds=0) + paused = _claim(tmp_path, scope=_SCOPE, pid=os.getpid(), age_seconds=two_hours) + other_namespace = _claim(tmp_path, scope=_OTHER_NAMESPACE, pid=_dead_pid(), age_seconds=two_hours) with run_write_lock(tmp_path): pass assert not abandoned.exists() + assert just_left.exists() and paused.exists() and other_namespace.exists() + + +def test_where_the_pid_scope_is_unknown_no_claim_is_removed(tmp_path, monkeypatch): + monkeypatch.setattr(storage, "_pid_scope", lambda: None) + abandoned = _claim(tmp_path, scope=_SCOPE, pid=_dead_pid(), age_seconds=2 * 3600) + + with run_write_lock(tmp_path): + pass + + assert abandoned.exists() + + +def test_a_new_claim_records_its_pid_scope_and_process_when_known(tmp_path, monkeypatch): + seen = [] + real_rename = os.rename + + def spy(src, dst): + seen.append(Path(src).name) + return real_rename(src, dst) + + monkeypatch.setattr(os, "rename", spy) + monkeypatch.setattr(storage, "_pid_scope", lambda: _SCOPE) + _publish(tmp_path, "a", now=1_700_000_000.0) + monkeypatch.setattr(storage, "_pid_scope", lambda: None) + _publish(tmp_path, "b", now=1_700_000_100.0) + + known, unknown = seen + assert storage._claim_owner(known) == (_SCOPE, os.getpid()) + assert storage._claim_owner(unknown) is None + + +def test_a_part_file_appears_whole_or_not_at_all(tmp_path, monkeypatch): + """A reader never meets a part file still being written, and a failed write leaves nothing behind.""" + import pyarrow.parquet as pq + + real_write = pq.write_table + + def half_written(table, where, **kwargs): + Path(where).write_bytes(b"PAR1 not yet a parquet file") + raise OSError("disk full") + + monkeypatch.setattr(pq, "write_table", half_written) + with pytest.raises(OSError, match="disk full"): + write_locale_partition(_rows("a"), tmp_path, locale="en_US", run_id="1700000000") + assert [p for p in tmp_path.rglob("*") if p.is_file()] == [] + + monkeypatch.setattr(pq, "write_table", real_write) + previous = os.umask(0o022) + try: + write_locale_partition(_rows("b"), tmp_path, locale="en_US", run_id="1700000000") + finally: + os.umask(previous) + (part,) = [p for p in tmp_path.rglob("*") if p.is_file()] + assert not part.name.startswith(".") and part.stat().st_mode & 0o777 == 0o644 + assert existing_trajectory_ids(tmp_path, run="1700000000") == {"b"} def test_the_lock_excludes_another_process(tmp_path): @@ -257,18 +351,55 @@ def hold(): assert waited > 0.5 -def test_where_flock_is_unsupported_writing_is_refused(tmp_path, monkeypatch): - """No fallback lock: ``filelock``'s can each let two writers in, so a root without ``flock`` takes no runs.""" +def _without_flock(monkeypatch, code: int = errno.ENOSYS) -> None: + """Make ``flock`` fail as it does on a filesystem without usable locks.""" fcntl = pytest.importorskip("fcntl") - import errno def unsupported(*args, **kwargs): - raise OSError(errno.ENOSYS, "flock is not supported here") + raise OSError(code, os.strerror(code)) monkeypatch.setattr(fcntl, "flock", unsupported) - with pytest.raises(RunLockUnsupportedError, match="not supported"): - check_run_write_lock(tmp_path) - with pytest.raises(RunLockUnsupportedError), run_write_lock(tmp_path): - pass - assert list(tmp_path.iterdir()) == [] + +@pytest.mark.parametrize( + "code", + [errno.ENOSYS, errno.ENOLCK, errno.EOPNOTSUPP, errno.ENOTSUP], + ids=["ENOSYS", "ENOLCK", "EOPNOTSUPP", "ENOTSUP"], +) +def test_where_flock_is_unsupported_runs_are_still_written(tmp_path, monkeypatch, caplog, recwarn, code): + """Written without the shared lock, with one warning per root, and never through a soft lock. + + Lustre without ``flock`` fails with ``ENOSYS``, NFS without its lock + service with ``ENOLCK``, and some network mounts with ``EOPNOTSUPP``. + """ + _without_flock(monkeypatch, code) + + with caplog.at_level(logging.WARNING, logger="usersim.engine"): + run_id, _ = save_sampled_rows(_rows("a"), tmp_path, locale="en_US", run_id=None, started_at=0.0) + save_sampled_rows(_rows("b"), tmp_path, locale="en_US", run_id=run_id, started_at=0.0) + + assert existing_trajectory_ids(tmp_path, run=run_id) == {"a", "b"} + assert sum("not supported" in record.getMessage() for record in caplog.records) == 1 + assert not [warning for warning in recwarn if "SoftFileLock" in str(warning.message)] + + +def test_where_flock_is_unsupported_the_threads_of_a_process_still_take_turns(tmp_path, monkeypatch): + _without_flock(monkeypatch) + held = threading.Event() + + def hold(): + with run_write_lock(tmp_path): + held.set() + time.sleep(1.0) + + holder = threading.Thread(target=hold) + holder.start() + try: + assert held.wait(timeout=30) + started = time.monotonic() + with run_write_lock(tmp_path): + waited = time.monotonic() - started + finally: + holder.join(timeout=30) + + assert waited > 0.5 diff --git a/tests/test_cli_simulate_run_partition.py b/tests/test_cli_simulate_run_partition.py index 0bea142..6f9e29d 100644 --- a/tests/test_cli_simulate_run_partition.py +++ b/tests/test_cli_simulate_run_partition.py @@ -458,36 +458,48 @@ def create(self, builder, num_records=None, artifact_path=None): assert tmp_path not in folders[0].parents -def test_a_root_without_file_locks_is_refused_before_any_model_call(tmp_path: Path, monkeypatch) -> None: - """The user is told which folder and why, as a configuration error, before anything is simulated.""" +def test_a_root_without_file_locks_still_gets_its_run_with_one_warning(tmp_path: Path, monkeypatch, caplog) -> None: + """Shared filesystems such as Lustre are often mounted without ``flock``; a run is written there anyway.""" fcntl = pytest.importorskip("fcntl") import errno + import logging from usersim.cli import simulate - from usersim.cli._errors import ConfigError from usersim.cli._models import default_models_path, load_models_config def unsupported(*args, **kwargs): raise OSError(errno.ENOSYS, "flock is not supported here") - def no_pipeline(**kwargs): - raise AssertionError("a pipeline was built") + class _FakeResult: + def load_dataset(self): + return _synthetic_trajectory("FRESH") + + class _FakeDD: + def create(self, builder, num_records=None, artifact_path=None): + return _FakeResult() monkeypatch.setattr(fcntl, "flock", unsupported) + monkeypatch.chdir(tmp_path) + out = tmp_path / "trajs" with ( - patch("usersim.cli._pipeline.build_simulator_config_builder", no_pipeline), - pytest.raises(ConfigError, match="file locks"), + patch("usersim.cli._pipeline.build_simulator_config_builder", lambda **kwargs: (_FakeDD(), object())), + caplog.at_level(logging.WARNING, logger="usersim.engine"), ): - simulate._simulate_to_parquet( + rc = simulate._simulate_to_parquet( models=load_models_config(default_models_path()), models_path=default_models_path(), panel=None, locales=["en_US"], num_rows=1, - out=tmp_path / "trajs", + out=out, assets_dir=packaged_assets_dir(), max_turns=1, probe_mix={"general_open_ended": 1.0}, random_seed=42, skip_existing=True, ) + + assert rc == 0 + (run_id,) = list_runs(out) + assert set(read_partitioned_dataset(out, run=run_id)["trajectory_id"].astype(str)) == {"FRESH"} + assert sum("not supported" in record.getMessage() for record in caplog.records) == 1 diff --git a/tests/test_runs.py b/tests/test_runs.py index e5e0a7e..cee3130 100644 --- a/tests/test_runs.py +++ b/tests/test_runs.py @@ -303,8 +303,6 @@ def test_simulate_notebook_writes_runs_as_the_cli_does() -> None: a run with another writer started in the same second, or save an id twice. """ src = _simulate_notebook_source() - # Refused before any model call where the trajectory root has no file locks. - assert "check_run_write_lock(TRAJECTORY_PATH)" in src assert "sampled_run_to_resume(TRAJECTORY_PATH)" in src assert "save_sampled_rows(" in src assert "write_locale_partition(" not in src and "new_run_id(" not in src