diff --git a/notebooks/01_simulate.ipynb b/notebooks/01_simulate.ipynb index 776b740..3071a01 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,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", @@ -577,12 +581,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." ] @@ -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", @@ -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", @@ -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}\")" ] }, { diff --git a/pyproject.toml b/pyproject.toml index ecdb97c..32611b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/src/usersim/cli/simulate.py b/src/usersim/cli/simulate.py index 71a144b..dc520e7 100644 --- a/src/usersim/cli/simulate.py +++ b/src/usersim/cli/simulate.py @@ -27,6 +27,8 @@ import argparse import logging +import tempfile +import time from fractions import Fraction from pathlib import Path @@ -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", @@ -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: @@ -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, @@ -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, diff --git a/src/usersim/engine/core/storage.py b/src/usersim/engine/core/storage.py index 64dcb88..536820c 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,19 @@ from __future__ import annotations +import errno +import json import logging +import os +import re import shutil import tempfile +import threading 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 +198,331 @@ 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 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 +#: 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-" + +#: 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() + +#: 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: + """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 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. 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 = 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 _pid_scope() -> str | None: + """The scope in which this process's PIDs identify processes: this boot and this PID namespace. + + 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. + """ + 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: + """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 not in _NO_FLOCK_ERRNOS: + 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 / _claim_name() + 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 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. + """ + 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, @@ -401,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 @@ -996,20 +1340,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..ed7b410 --- /dev/null +++ b/tests/engine/core/test_run_writes.py @@ -0,0 +1,405 @@ +# 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 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, + 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, + write_locale_partition, +) + +#: 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-*")) + + +#: 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): + """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 _without_flock(monkeypatch, code: int = errno.ENOSYS) -> None: + """Make ``flock`` fail as it does on a filesystem without usable locks.""" + fcntl = pytest.importorskip("fcntl") + + def unsupported(*args, **kwargs): + raise OSError(code, os.strerror(code)) + + monkeypatch.setattr(fcntl, "flock", unsupported) + + +@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 8da7ab9..6f9e29d 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,98 @@ 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_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._models import default_models_path, load_models_config + + def unsupported(*args, **kwargs): + raise OSError(errno.ENOSYS, "flock is not supported here") + + 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", lambda **kwargs: (_FakeDD(), object())), + caplog.at_level(logging.WARNING, logger="usersim.engine"), + ): + 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 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 7d54230..cee3130 100644 --- a/tests/test_runs.py +++ b/tests/test_runs.py @@ -296,6 +296,22 @@ 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() + 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 000edad..64bb871 100644 --- a/uv.lock +++ b/uv.lock @@ -4649,6 +4649,7 @@ version = "0.2.0.dev0" source = { editable = "." } dependencies = [ { name = "data-designer" }, + { name = "filelock" }, { name = "httpx" }, { name = "huggingface-hub" }, { name = "langdetect" }, @@ -4685,6 +4686,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "data-designer", specifier = ">=0.9.4,<0.10" }, + { name = "filelock" }, { name = "httpx" }, { name = "huggingface-hub" }, { name = "langdetect" },