From 1563fbf5c7f6e0e376b8de00a0ef56518506a402 Mon Sep 17 00:00:00 2001 From: nichinichisou Date: Sun, 18 Oct 2026 19:45:00 +0800 Subject: [PATCH] feat(depot): release completed bundle downloads during exports --- CHANGELOG.md | 1 + docs/architecture.md | 5 ++ docs/commands.md | 8 ++ docs/configuration.md | 9 ++- src/nnnotes/cli/assets.py | 10 +++ src/nnnotes/core/downloads.py | 97 ++++++++++++++++------- src/nnnotes/pipeline/downloads.py | 56 ++++++++++++++ src/nnnotes/pipeline/orchestrate.py | 34 ++++++++- tests/cli/test_assets.py | 34 ++++++++- tests/core/test_downloads.py | 65 ++++++++++++++++ tests/pipeline/test_downloads.py | 104 +++++++++++++++++++++++++ tests/pipeline/test_orchestrate.py | 114 +++++++++++++++++++++++++++- 12 files changed, 498 insertions(+), 39 deletions(-) create mode 100644 src/nnnotes/pipeline/downloads.py create mode 100644 tests/pipeline/test_downloads.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 423fb79..ac9277d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ schemas in `nnnotes.schemas` and the documented Python API. Before 1.0 a minor v ### Features - Store content, task inputs, results and catalog versions in a persistent depot with optional remote storage, verified collection and bounded source downloads. +- Release recoverable bundle downloads after successful conversion batches under the configured budget and command leases. - Assemble Live2D, room, spot, live, story and UI products from shared object tables and explicit manifest references. - Generate localized live and story text products with game or supplied fonts, shared glyph atlases and material records. - Build named sites through depot product, player read-set, asset and index tasks. diff --git a/docs/architecture.md b/docs/architecture.md index d3e28b9..7fb53b9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -48,6 +48,11 @@ The site's `Builder` composes product snapshots, catalog indexes and player page Converters use bounded process-local memoization for repeated reads and computations. Persistent reuse is keyed by depot tasks and their content inputs. Clearing a process memo releases its values; a later task can reuse its committed depot result. +The coordinator observes successful object batches after worker readers close. Their recoverable bundle downloads +can then be released toward the configured budget under the command's exclusive lease. Preliminary episode +requests retain inputs until the expanded product request is known. Stored conversion results remain available +for assembly and layouts. + `export episode`, `export shader`, `export audio` and `music-data` expose subject-specific data and conversions through the Session. All commands use the readers and converters below the pipeline layer. ## Outputs diff --git a/docs/commands.md b/docs/commands.md index 87f2979..6969eef 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -69,6 +69,14 @@ Catalog files are kept. A removed file is downloaded again when a command needs identity (sha256 and size) of every bundle and raw file it reads, so `export assets` plans with files whose name carries the catalog's hash without downloading them again. +During product and asset requests, a completed object-conversion batch makes its bundle downloads eligible for +release toward the same budget. Readers finish before release; active batches, failed batches and preliminary +episode requests stay protected. A batch already present in the depot can release its downloaded input as well. +Only files recoverable through an explicit catalog or depot locator are released, and their contents are checked +against the input identity. `all` retains downloads throughout the run; `0` releases each eligible batch without +scanning the download directory. With a finite size, finished files are evicted in LRU order and active files can +temporarily exceed the budget. Later requests can download a released bundle again. + Commands may run at the same time on one depot: no command removes downloads while another one holds a lease, and a command that starts while one trims waits for it. The lease of a process that ended without releasing it is dropped by the next command; a lease written on another host counts for a day. diff --git a/docs/configuration.md b/docs/configuration.md index f36ad31..a5419a3 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -110,7 +110,7 @@ exist when they are given. | `[paths] player` | `NNNOTES_PATHS_PLAYER` | `--player` (of `site build`) | ournotes-player checkout (built) or installed package | | `[paths] fonts.` | `NNNOTES_PATHS_FONTS_` (e.g. `NNNOTES_PATHS_FONTS_ZH_HANT`) | `--font =PATH` (of `site build`) | font file (OpenType or TrueType) the story text of a language is drawn with (`--fonts open`); also read for the other story languages whose game fonts fall back to that language's font | | `[paths] fonts.emoji` | `NNNOTES_PATHS_FONTS_EMOJI` | `--font emoji=PATH` (of `site build`) | colour emoji font (PNG bitmap glyphs: CBDT or sbix) the stories' emoji sprites are drawn from (`--fonts open`); Noto Color Emoji (SIL Open Font License 1.1) is the tested font. Optional: without it the sprites keep their layout with empty glyphs, listed in `ui/fonts.json` `coverage.sprites.missing` | -| `[depot] downloads` | `NNNOTES_DEPOT_DOWNLOADS` | — | how much of the depot's downloaded bundles and raw files to keep once no command uses the depot: a size (`500MB`, `20GiB`; units B, KB, MB, GB, TB and KiB, MiB, GiB, TiB), `all`, or `0` (unset) ([commands.md](commands.md#downloads)) | +| `[depot] downloads` | `NNNOTES_DEPOT_DOWNLOADS` | — | download-cache budget, applied to finished conversion batches and when the last command ends: a size (`500MB`, `20GiB`; units B, KB, MB, GB, TB and KiB, MiB, GiB, TiB), `all`, or `0` (unset) ([commands.md](commands.md#downloads)) | | `[depot] remote` | `NNNOTES_DEPOT_REMOTE` | — | a remote of the depot shared by the machines that use it: `s3:///` (S3-compatible object storage; `pip install 'nnnotes[remote]'`) or `file:///` ([contracts.md](contracts.md#depot-remote)) | | `[depot] endpoint` | `NNNOTES_DEPOT_ENDPOINT` | — | endpoint URL of the S3-compatible object storage; unset: AWS S3 | | `[depot] access_key_id` | `NNNOTES_DEPOT_ACCESS_KEY_ID` | — | access key id of the object storage; unset: boto3 looks for credentials itself (`AWS_ACCESS_KEY_ID`, `~/.aws`) | @@ -188,9 +188,10 @@ story text); it is not a setting. | `catalog versions` / `import` / `diff`, `depot verify` | the depot; `import` reads the APK's catalog from `[paths] apk` when it is set | | `catalog fetch` | the depot, `[catalog] region` and `language`, that region's `cdn`; its `api` (optional) for the resource version label | -The region, its `cdn`, `[bundle] key` and `nonce_seed` are read only when a file must be downloaded: when the -catalog and every bundle a command needs are downloaded, none of them is needed. A download that needs a missing -one stops the command with one line naming the file and the setting (see Errors). +The region, its `cdn`, `[bundle] key` and `nonce_seed` are read only when a file must be downloaded. A product +request can read a downloaded catalog and its bundle files without these settings. Use `[depot] downloads = all` +for an offline sequence of requests; a bounded budget can release a bundle that a later request needs again. +A download that needs a missing setting stops the command with one line naming the file and setting (see Errors). ## Errors diff --git a/src/nnnotes/cli/assets.py b/src/nnnotes/cli/assets.py index d77d215..d36ca5a 100644 --- a/src/nnnotes/cli/assets.py +++ b/src/nnnotes/cli/assets.py @@ -1317,8 +1317,10 @@ def cmd_export(args, cfg): def request_orchestrator(ws, cfg): """The shared executor for product requests.""" + from ..pipeline.downloads import BundleDownloads from ..pipeline.orchestrate import Orchestrator + lease, budget = getattr(ws.args, "lease", None), Session(cfg).downloads_budget() return Orchestrator( ws.depot, ws.scheduled, @@ -1329,12 +1331,14 @@ def request_orchestrator(ws, cfg): initializer=_worker_init, initargs=(cfg,), log=ws.log, + on_batch_done=BundleDownloads(ws.depot, lease, budget) if lease is not None and budget is not None else None, ) def run_request(args, cfg, ws: Workspace, expand=None, prepared=None): """Execute a product request and optional expansion without materializing an output layout.""" from ..pipeline import builtins + from ..pipeline.downloads import BundleDownloads from ..pipeline.orchestrate import failure if prepared is None: @@ -1343,6 +1347,8 @@ def run_request(args, cfg, ws: Workspace, expand=None, prepared=None): orch = request_orchestrator(ws, cfg) else: pipeline, params, orch = prepared + if expand is not None: + orch.on_batch_done = None # the episode supplies the roots of the complete product request ws.run_index() facts = ws.facts(pipeline, fetch=True, background=True) sel = ws.selected @@ -1374,6 +1380,10 @@ def run_request(args, cfg, ws: Workspace, expand=None, prepared=None): aborted = aborted or str(e) if isinstance(facts["raw"], Later) and aborted is None: ws.log(f"{sum(p['kind'] == 'raw' for p in ws.problems())} raw files not readable") + released = orch.on_batch_done + if isinstance(released, BundleDownloads) and released.removed: + run.log.append({"event": "downloads.released", "files": released.removed, "bytes": released.bytes_removed}) + ws.log(f"downloads: released {released.removed} bundle files ({released.bytes_removed} bytes)") if expand is not None and aborted is None: more = expand(run) if more is not None: diff --git a/src/nnnotes/core/downloads.py b/src/nnnotes/core/downloads.py index 732afcc..f1a79a1 100644 --- a/src/nnnotes/core/downloads.py +++ b/src/nnnotes/core/downloads.py @@ -1,11 +1,12 @@ -"""The downloads of a depot as a cache bounded in size: what a command reads again stays while commands run, and the -bundles and raw files are trimmed to the budget when the last command using the depot ends. +"""Budgeted depot downloads, with finished-batch release and command leases protecting active readers. /downloads/ the game files as downloaded (source.catalog) /leases/.json one per running command: {"host", "pid", "started"} /leases/.exclusive held by a command that must be the only one (trimming, `depot gc`) -A command holds a lease while it runs (`lease`). When it ends and no other lease is live (a lease of this host whose +A command holds a lease while it runs (`lease`). A successful object conversion batch may release its recoverable +bundle downloads toward the budget (`release`) while it is the only command. When the command ends and no other +lease is live (a lease of this host whose process runs, or of another host younger than a day), the files under the `bundles/` and `raw/` directories of downloads/ are removed, least recently used first (a file's modification time: the catalog sets it when it hands the file out), until they fit the budget. Catalog files stay. A removed file is downloaded again when a command needs @@ -25,12 +26,12 @@ import secrets import socket import time -from collections.abc import Generator, Iterator +from collections.abc import Generator, Iterable, Iterator from contextlib import contextmanager from pathlib import Path from typing import Any -__all__ = ["evictable", "exclusive", "lease", "parse_budget", "trim"] +__all__ = ["evictable", "exclusive", "lease", "parse_budget", "release", "trim"] FOREIGN_LEASE_SECONDS = 24 * 3600 # a lease of another host counts as live this long EXCLUSIVE = ".exclusive" @@ -71,39 +72,79 @@ def evictable(downloads: Path) -> Iterator[Path]: yield p -def trim(downloads: Path, budget: int | None, *, dry_run: bool = False) -> dict[str, int]: - """Remove the least recently used bundles and raw files until the rest fit `budget` bytes (None: nothing). - -> {"kept", "removed", "bytesKept", "bytesRemoved"}.""" - files: list[tuple[int, int, Path]] = [] - for p in evictable(downloads): +def _files(paths: Iterable[Path]) -> list[tuple[int, int, Path]]: + out: list[tuple[int, int, Path]] = [] + for path in paths: try: - st = p.stat() + stat = path.stat() except OSError: continue - files.append((st.st_mtime_ns, st.st_size, p)) + out.append((stat.st_mtime_ns, stat.st_size, path)) + return out + + +def _remove(files: list[tuple[int, int, Path]], needed: int, dry_run: bool) -> tuple[list[Path], int]: + removed: list[Path] = [] + removed_bytes = 0 + for _, size, path in sorted(files, key=lambda f: (f[0], str(f[2]))): + if removed_bytes >= needed: + break + if not dry_run: + try: + path.unlink() + except OSError: + continue + removed.append(path) + removed_bytes += size + return removed, removed_bytes + + +def trim(downloads: Path, budget: int | None, *, dry_run: bool = False) -> dict[str, int]: + """Remove the least recently used bundles and raw files until the rest fit `budget` bytes (None: nothing). + -> {"kept", "removed", "bytesKept", "bytesRemoved"}.""" + files = _files(evictable(downloads)) total = sum(f[1] for f in files) - removed = removed_bytes = 0 - if budget is not None: - for _, size, p in sorted(files, key=lambda f: (f[0], str(f[2]))): - if total - removed_bytes <= budget: - break - if not dry_run: - try: - p.unlink() - except OSError: # open elsewhere (Windows): kept - continue - removed += 1 - removed_bytes += size - if not dry_run and removed: - _prune_dirs(downloads) + removed, removed_bytes = _remove(files, 0 if budget is None else total - budget, dry_run) + if not dry_run and removed: + _prune_dirs(downloads) return { - "kept": len(files) - removed, - "removed": removed, + "kept": len(files) - len(removed), + "removed": len(removed), "bytesKept": total - removed_bytes, "bytesRemoved": removed_bytes, } +def release(downloads: Path, ready: Iterable[Path], budget: int | None) -> dict[str, int]: + """Trim only finished downloads, in LRU order, toward the cache budget. The caller holds an exclusive lease. + + Active files remain protected even when they exceed the budget. A zero budget examines only finished files. + Paths must resolve below a bundles/ or raw/ directory within downloads; partial files are retained. + """ + if budget is None: + return {"removed": 0, "bytesRemoved": 0} + root = downloads.resolve() + paths: set[Path] = set() + for path in ready: + path = path.resolve() + if not path.is_relative_to(root) or not {"bundles", "raw"} & set(path.relative_to(root).parts[:-1]): + raise ValueError("released file must be a bundle or raw file within the downloads directory") + if path.is_file() and not path.name.endswith(".part"): + paths.add(path) + files = _files(paths) + total = sum(f[1] for f in files) if budget == 0 else sum(f[1] for f in _files(evictable(root))) + removed, size = _remove(files, total - budget, False) + for path in removed: + parent = path.parent + while parent != root: + try: + parent.rmdir() + except OSError: + break + parent = parent.parent + return {"removed": len(removed), "bytesRemoved": size} + + def _prune_dirs(root: Path) -> None: for d in sorted((p for p in root.rglob("*") if p.is_dir()), key=lambda p: len(p.parts), reverse=True): try: diff --git a/src/nnnotes/pipeline/downloads.py b/src/nnnotes/pipeline/downloads.py new file mode 100644 index 0000000..24d4397 --- /dev/null +++ b/src/nnnotes/pipeline/downloads.py @@ -0,0 +1,56 @@ +"""Budgeted release of recoverable bundle downloads after complete conversion batches.""" + +from pathlib import Path + +from ..core import downloads +from .contract import Input + + +class BundleDownloads: + """A main-process batch observer; only closed, recoverable bundle inputs become eligible for release.""" + + def __init__(self, depot, lease: Path | None, budget: int | None): + self.depot, self.lease, self.budget = depot, Path(lease) if lease is not None else None, budget + self.ready: dict[Path, Input] = {} + self.removed = self.bytes_removed = 0 + + def __call__(self, tasks): + if self.budget is None or self.lease is None or not self.lease.is_file(): + return + if self.lease.parent.resolve() != (self.depot.root / "leases").resolve() or self.lease.suffix != ".json": + return + if not tasks or any(task.stage != "unity.export" for task in tasks): + return + root = self.depot.downloads.resolve() + for task in tasks: + inp = task.input("bundle") + kinds = {loc["kind"] for loc in inp.locators} + if "catalog" not in kinds and not ("depot" in kinds and self.depot.has(inp.sha256, inp.size)): + continue + for loc in inp.locators: + if loc["kind"] != "download": + continue + path = (root / loc["path"]).resolve() + if path.is_relative_to(root) and "bundles" in path.relative_to(root).parts[:-1]: + self.ready[path] = inp + if not self.ready: + return + with downloads.exclusive(self.depot.root, self.lease) as alone: + if not alone: + return + valid = [] + for path, inp in list(self.ready.items()): + try: + matches = path.is_file() and self.depot.identify(path) == (inp.sha256, inp.size) + except OSError: + matches = False + if matches: + valid.append(path) + else: + self.ready.pop(path) + result = downloads.release(root, valid, self.budget) + self.removed += result["removed"] + self.bytes_removed += result["bytesRemoved"] + for path in valid: + if not path.exists(): + self.ready.pop(path, None) diff --git a/src/nnnotes/pipeline/orchestrate.py b/src/nnnotes/pipeline/orchestrate.py index 8006c1b..dd48d8c 100644 --- a/src/nnnotes/pipeline/orchestrate.py +++ b/src/nnnotes/pipeline/orchestrate.py @@ -510,7 +510,9 @@ def coverage(results, census: dict | None = None) -> dict: class Orchestrator: """Runs pipelines of stages into a depot (see the module documentation). `initializer(depot, *initargs)` runs - first in every worker (settings, caches, the depot's fetcher); it must be importable (spawned processes).""" + first in every worker (settings, caches, the depot's fetcher); it must be importable (spawned processes). + `on_batch_done(tasks)` observes successful batches in the coordinator after readers close, including batches + entirely satisfied by committed results. Failed batches and incomplete task descriptions are not observed.""" def __init__( self, @@ -525,12 +527,16 @@ def __init__( initializer=None, initargs: tuple = (), log=_log, + on_batch_done=None, ): self.depot, self.stages = depot, stages self.workers, self.memory, self.worker_bytes = max(0, int(workers)), memory, int(worker_bytes) self.recycle_tasks, self.recycle_rss = recycle_tasks, recycle_rss self.initializer, self.initargs = initializer, tuple(initargs) self.log = log + self.on_batch_done = on_batch_done + self._batch_errors: set[tuple] = set() + self._observe_batches = True self._pool: list[_Worker] = [] self._started = 0 @@ -595,7 +601,7 @@ def _stage(self, stage: Stage, params: dict, env: Env, force: set, run: Run) -> t0 = time.perf_counter() tasks: list[Task] = [] try: - subjects = stage.subjects(env) + subjects = list(stage.subjects(env)) except Pending as p: self._fail( run, @@ -623,6 +629,7 @@ def _stage(self, stage: Stage, params: dict, env: Env, force: set, run: Run) -> ) except Exception as e: self._fail(run, env, tid, None, stage.name, "task.describe", f"{type(e).__name__}: {e}"[:2000]) + self._observe_batches = len(tasks) == len(subjects) forced = "*" in force or stage.name in force todo = [] for t in sorted(tasks, key=lambda t: t.id): @@ -631,6 +638,11 @@ def _stage(self, stage: Stage, params: dict, env: Env, force: set, run: Run) -> env.done(t.id, t.key) else: todo.append(t) + pending = {t.id for t in todo} + if self.on_batch_done is not None and self._observe_batches: + for job in jobs(tasks, self.stages): + if job.batch is not None and not any(t.id in pending for t in job.tasks): + self.on_batch_done(job.tasks) replies = self._execute(todo, force=lambda t: forced or t.id in force, run=run) stop = next((r for r in replies.values() if r.get("code") == "config"), None) for t in todo: @@ -657,6 +669,7 @@ def _fail(run: Run, env: Env, tid: str, key, stage: str, code: str, message: str # ---------------------------------------------------------------- execution def _execute(self, tasks: list[Task], force, run: Run) -> dict[str, dict]: """Run `tasks`; {task id: reply} (none for a task not started because the run stopped).""" + self._batch_errors.clear() if not tasks: return {} sched = Scheduler(jobs(tasks, self.stages), self.task_budget()) @@ -675,6 +688,7 @@ def _execute(self, tasks: list[Task], force, run: Run) -> dict[str, dict]: break if j.batch is not None: self.stages[j.batch[0]].end_batch(j.batch[1]) + self._batch_done(j, replies) sched.finished(j.id) else: self._pool_run(sched, force, replies, run) @@ -724,6 +738,7 @@ def _collect(self, w: _Worker, sched: Scheduler, replies: dict, run: Run) -> Non except (EOFError, OSError): break if msg.get("end"): + self._batch_done(job, replies) sched.finished(job.id) w.job = None w.done += w.replied @@ -733,10 +748,14 @@ def _collect(self, w: _Worker, sched: Scheduler, replies: dict, run: Run) -> Non t = job.tasks[w.replied] w.replied += 1 replies[t.id] = msg + if not msg["ok"] and job.batch is not None: + self._batch_errors.add(job.batch) self._done(run, t, msg, w) if msg.get("code") == "config": sched.stop() w.proc.join(5) + if job.batch is not None: + self._batch_errors.add(job.batch) if w.replied < len(job.tasks): t = job.tasks[w.replied] replies[t.id] = { @@ -754,6 +773,17 @@ def _collect(self, w: _Worker, sched: Scheduler, replies: dict, run: Run) -> Non w.job = None self._replace(w) + def _batch_done(self, job: Job, replies: dict) -> None: + """Observe a successful batch after its reader is dropped; cache-only batches are also observed.""" + if ( + self.on_batch_done is not None + and self._observe_batches + and job.batch is not None + and job.batch not in self._batch_errors + and all(replies.get(t.id, {}).get("ok") for t in job.tasks) + ): + self.on_batch_done(job.tasks) + def _recycle(self, w: _Worker, r: int | None) -> bool: if self.recycle_tasks and w.done >= self.recycle_tasks: return True diff --git a/tests/cli/test_assets.py b/tests/cli/test_assets.py index cc9e3eb..e3e29f1 100644 --- a/tests/cli/test_assets.py +++ b/tests/cli/test_assets.py @@ -214,19 +214,24 @@ def depends(self, subject, env): return [self._census(subject, env), "link.scripts:all"] def inputs(self, subject, env): - return [env.input_of(self._census(subject, env), "census"), env.input_of("link.scripts:all", "scripts")] + return [ + env.fact("bundles")[units.table(env).unit(subject).bundle], + env.input_of(self._census(subject, env), "census"), + env.input_of("link.scripts:all", "scripts"), + ] def context(self, subject, env): unit = units.table(env).unit(subject) return {"unit": unit.root, "members": list(unit.members)} def batch(self, task): - return task.input("census").sha256 + return task.input("bundle").sha256 def estimate(self, subject, env, inputs): return Cost(0.01, 1 << 20) def run(self, task, depot): + assert depot.input_bytes(task.input("bundle")).startswith(b"UnityFS") doc = contract.loads(depot.input_bytes(task.input("census"))) scripts = contract.loads(depot.input_bytes(task.input("scripts")))["scripts"] container = {(e["file"], e["pathId"]): e["path"] for b in doc["assetBundles"] for e in b["container"]} @@ -1586,10 +1591,12 @@ def test_export_live2d_writes_the_cubism_files_of_one_model(tmp_path, capsys, mo assert f"live2d.cubism:{MODEL_KEY}" in ids and "fake.product:Bg/x" not in ids -def test_export_runs_the_request_expand_gives_and_places_the_last_run(tmp_path, capsys, monkeypatch): +@pytest.mark.parametrize("downloads_budget", ["all", "0"]) +def test_export_runs_the_request_expand_gives_and_places_the_last_run(tmp_path, capsys, monkeypatch, downloads_budget): from urllib.parse import quote monkeypatch.setattr(assets, "STAGE_SOURCES", PRODUCT_SOURCES) + monkeypatch.setenv("NNNOTES_DEPOT_DOWNLOADS", downloads_budget) monkeypatch.setattr("nnnotes.pipeline.catalogdb.apk_version", lambda apk: None) d = setup_data(tmp_path, censuses(broken=False), extra=[(MODEL_KEY, "Assets/Game/Bg/x.prefab", [3])]) apk = builtin_apk(tmp_path / "base.apk", {"assets/aa/catalog.bin": synth.CatalogWriter().build([])}) @@ -1601,6 +1608,7 @@ def hook(run): # called once, with the first run (the command's own request: th if len(calls) > 1: pytest.fail("expand was called again after its second answer") assert "fake.product:Bg/x" in calls[0] + assert (d["downloads"] / "bundles" / NAMES["bg_x"]).exists() return {"live2d.cubism": [MODEL_KEY]}, [f"key:{MODEL_KEY}"] def hooked(args, cfg, ws, expand=None): # the command's export takes the hook; its own re-run takes none @@ -1634,6 +1642,26 @@ def hooked(args, cfg, ws, expand=None): # the command's export takes the hook; assert not (tmp_path / "out" / "fake.product").exists() # the first run's product is not placed +def test_completed_bundle_downloads_are_released_before_layout_publication(tmp_path, capsys, fake, monkeypatch): + from nnnotes.pipeline import layout + + monkeypatch.setenv("NNNOTES_DEPOT_DOWNLOADS", "0") + d = setup_data(tmp_path, censuses(broken=False)) + place = layout.materialize + observed = [] + + def publish(*args, **kwargs): + assert list((tmp_path / "s/leases").glob("*.json")) + assert not (d["downloads"] / "bundles" / NAMES["bg_x"]).exists() + observed.append(True) + return place(*args, **kwargs) + + monkeypatch.setattr(layout, "materialize", publish) + code, out, err = export(capsys, tmp_path, d, "s", "out") + assert code == 0 and observed, err + assert "downloads: released" in err + + def test_every_product_a_story_plays_with_is_installed(): from nnnotes.pipeline import story diff --git a/tests/core/test_downloads.py b/tests/core/test_downloads.py index 656645a..e14751b 100644 --- a/tests/core/test_downloads.py +++ b/tests/core/test_downloads.py @@ -137,3 +137,68 @@ def test_exclusive_is_granted_to_the_only_command(tmp_path): with downloads.exclusive(tmp_path, own) as alone: assert not alone # another lease is live assert not list((tmp_path / "leases").iterdir()) + + +def test_release_respects_the_budget_and_protects_active_files(tmp_path): + root = tmp_path / "downloads" + old = put(root, "bundles/old.bundle", 10, 1) + ready = put(root, "bundles/ready.bundle", 10, 2) + active = put(root, "bundles/active.bundle", 20, 3) + assert downloads.release(root, [old, ready], 30) == {"removed": 1, "bytesRemoved": 10} + assert not old.exists() and ready.exists() and active.exists() + assert downloads.release(root, [ready], 5) == {"removed": 1, "bytesRemoved": 10} + assert active.exists() + + +def test_zero_budget_release_examines_only_finished_paths(tmp_path, monkeypatch): + root = tmp_path / "downloads" + ready = put(root, "bundles/ready.bundle", 10, 1) + active = put(root, "raw/active.bin", 20, 2) + partial = put(root, "bundles/fetch.part", 5, 3) + + def scan(_): + raise AssertionError("zero-budget release does not need a directory inventory") + + monkeypatch.setattr(downloads, "evictable", scan) + assert downloads.release(root, [ready, ready, partial], 0) == {"removed": 1, "bytesRemoved": 10} + assert active.exists() and partial.exists() + + +def test_release_validates_paths_and_keeps_unlimited_downloads(tmp_path): + root = tmp_path / "downloads" + file = put(root, "bundles/one.bundle", 10, 1) + assert downloads.release(root, [file], None) == {"removed": 0, "bytesRemoved": 0} + for path in [tmp_path / "outside", root / "catalog.bin"]: + with pytest.raises(ValueError, match="within the downloads"): + downloads.release(root, [path], 0) + assert file.exists() + + +def test_release_keeps_locked_files_and_continues(tmp_path, monkeypatch): + root = tmp_path / "downloads" + locked = put(root, "bundles/locked.bundle", 10, 1) + other = put(root, "bundles/other.bundle", 20, 2) + unlink = type(locked).unlink + + def remove(path, *args, **kwargs): + if path == locked: + raise PermissionError("open") + return unlink(path, *args, **kwargs) + + monkeypatch.setattr(type(locked), "unlink", remove) + assert downloads.release(root, [locked, other], 0) == {"removed": 1, "bytesRemoved": 20} + assert locked.exists() and not other.exists() + + +def test_release_rejects_symlinks_to_outside_files(tmp_path): + root = tmp_path / "downloads" + target = put(tmp_path, "outside.bin", 10, 1) + link = root / "bundles/link.bundle" + link.parent.mkdir(parents=True) + try: + link.symlink_to(target) + except OSError: + pytest.skip("symlink creation is unavailable") + with pytest.raises(ValueError, match="within the downloads"): + downloads.release(root, [link], 0) + assert target.exists() and link.is_symlink() diff --git a/tests/pipeline/test_downloads.py b/tests/pipeline/test_downloads.py new file mode 100644 index 0000000..ce38c73 --- /dev/null +++ b/tests/pipeline/test_downloads.py @@ -0,0 +1,104 @@ +"""Bundle release uses completed task inputs, recoverable locators and command leases.""" + +import hashlib +from dataclasses import replace + +from nnnotes.core import downloads +from nnnotes.pipeline.contract import Input, Task +from nnnotes.pipeline.depot import Depot +from nnnotes.pipeline.downloads import BundleDownloads + + +def bundle(depot, data=b"bundle", name="one.bundle", recover=True): + path = depot.downloads / "bundles" / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + locators = [{"kind": "download", "path": "bundles/" + name}] + if recover: + locators.append({"kind": "catalog", "catalog": "example", "location": name}) + inp = Input("bundle", hashlib.sha256(data).hexdigest(), len(data), name, tuple(locators)) + return path, Task("unity.export", 1, name, inputs=(inp,)) + + +def test_other_leases_protect_files_until_a_later_batch(tmp_path): + depot = Depot(tmp_path) + path, task = bundle(depot) + with downloads.lease(tmp_path, None) as own: + release = BundleDownloads(depot, own, 0) + with downloads.lease(tmp_path, None): + release([task]) + assert path.exists() and release.removed == 0 + release([task]) + assert not path.exists() and release.removed == 1 and release.bytes_removed == 6 + + +def test_only_matching_recoverable_bundle_inputs_are_released(tmp_path): + depot = Depot(tmp_path) + changed, task = bundle(depot) + changed.write_bytes(b"another") + private, private_task = bundle(depot, name="private.bundle", recover=False) + other, other_task = bundle(depot, name="other.bundle") + with downloads.lease(tmp_path, None) as own: + release = BundleDownloads(depot, own, 0) + release([task, private_task]) + release([replace(other_task, stage="unity.census")]) + assert changed.exists() and private.exists() and other.exists() + release([other_task]) + assert not other.exists() and release.removed == 1 + + +def test_missing_lease_and_unlimited_budget_keep_files(tmp_path): + depot = Depot(tmp_path) + path, task = bundle(depot) + BundleDownloads(depot, None, 0)([task]) + BundleDownloads(depot, tmp_path / "absent", 0)([task]) + invalid = tmp_path / "outside.json" + invalid.write_text("{}") + BundleDownloads(depot, invalid, 0)([task]) + with downloads.lease(tmp_path, None) as own: + BundleDownloads(depot, own, None)([task]) + assert path.exists() + + +def test_cas_recovery_requires_an_explicit_depot_locator(tmp_path): + depot = Depot(tmp_path) + path, task = bundle(depot, recover=False) + inp = task.inputs[0] + depot.put(path.read_bytes()) + with downloads.lease(tmp_path, None) as own: + release = BundleDownloads(depot, own, 0) + release([task]) + assert path.exists() + task = replace(task, inputs=(replace(inp, locators=(*inp.locators, {"kind": "depot"})),)) + release([task]) + assert not path.exists() and depot.input_bytes(task.inputs[0]) == b"bundle" + + +def test_released_input_can_be_fetched_for_an_expanded_request(tmp_path): + fetched = [] + + def fetch(inp, locator): + fetched.append(locator) + path = tmp_path / "downloads/bundles/one.bundle" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"bundle") + return path + + depot = Depot(tmp_path, fetch=fetch) + path, task = bundle(depot) + with downloads.lease(tmp_path, None) as own: + BundleDownloads(depot, own, 0)([task]) + assert not path.exists() + assert depot.input_bytes(task.inputs[0]) == b"bundle" and len(fetched) == 1 + + +def test_download_locators_cannot_release_files_outside_the_depot(tmp_path): + depot = Depot(tmp_path / "depot") + outside = tmp_path / "bundles/outside" + outside.parent.mkdir() + outside.write_bytes(b"bundle") + _, task = bundle(depot) + inp = replace(task.inputs[0], locators=({"kind": "download", "path": str(outside)}, {"kind": "catalog"})) + with downloads.lease(depot.root, None) as own: + BundleDownloads(depot, own, 0)([replace(task, inputs=(inp,))]) + assert outside.exists() diff --git a/tests/pipeline/test_orchestrate.py b/tests/pipeline/test_orchestrate.py index 47c2914..4b3533e 100644 --- a/tests/pipeline/test_orchestrate.py +++ b/tests/pipeline/test_orchestrate.py @@ -172,11 +172,11 @@ def _note(self, line: str) -> None: f.write(line + "\n") -def batched(tmp_path, name, texts, workers): +def batched(tmp_path, name, texts, workers, **kwargs): log = tmp_path / f"{name}-log" log.mkdir() stages = registry(Batched(log), Count(), Total()) - depot, run, orch = orchestrate(tmp_path, name, texts, workers=workers, stages=stages) + depot, run, orch = orchestrate(tmp_path, name, texts, workers=workers, stages=stages, **kwargs) return depot, run, orch, [" ".join(f.read_text(encoding="utf-8").split()) for f in sorted(log.iterdir())] @@ -214,6 +214,116 @@ def test_a_worker_dying_in_a_job_fails_its_task_and_the_rest_runs_in_another(tmp assert orch._started >= 2 +@pytest.mark.parametrize("workers", [0, 2]) +def test_batch_observer_runs_after_readers_close_and_observes_cache_hits(tmp_path, workers): + seen = [] + log = tmp_path / "s-log" + + def observe(tasks): + batch = tasks[0].subject[0] + text = "\n".join(p.read_text(encoding="utf-8") for p in log.glob("*.log")) + assert f"drop {batch}" in text + seen.append([t.subject for t in tasks]) + + depot, first, _, logs = batched(tmp_path, "s", {"a1": "x", "a2": "y", "b1": "z"}, workers, on_batch_done=observe) + assert first.exit_code == 0 and sorted(seen) == [["a1", "a2"], ["b1"]] + seen.clear() + stages = registry(Batched(log), Count(), Total()) + again = Orchestrator(depot, stages, workers=workers, on_batch_done=observe, log=lambda _: None) + run = again.run(PIPELINE, facts={"files": toy_files(tmp_path / "in", {"a1": "x", "a2": "y", "b1": "z"})}) + assert all(t["status"] == "hit" for t in run.tasks.values()) + assert sorted(seen) == [["a1", "a2"], ["b1"]] + assert [" ".join(p.read_text(encoding="utf-8").split()) for p in sorted(log.iterdir())] == logs + + +@pytest.mark.parametrize("workers", [0, 1]) +def test_a_mixed_cached_batch_is_observed_after_its_pending_tasks(tmp_path, workers): + depot, _, _, _ = batched(tmp_path, "s", {"a1": "x"}, workers) + seen = [] + orch = Orchestrator( + depot, + registry(Batched(tmp_path / "s-log"), Count(), Total()), + workers=workers, + on_batch_done=lambda tasks: seen.append([t.subject for t in tasks]), + log=lambda _: None, + ) + run = orch.run(PIPELINE, facts={"files": toy_files(tmp_path / "in", {"a1": "x", "a2": "y"})}) + assert run.tasks["toy.src:a1"]["status"] == "hit" + assert run.tasks["toy.src:a2"]["status"] == "ran" + assert seen == [["a2"]] + + +@pytest.mark.parametrize("workers, failure", [(0, "CRASH"), (2, "CRASH"), (1, "DIE")]) +def test_failed_batches_are_kept_even_when_their_remaining_tasks_succeed(tmp_path, workers, failure): + seen = [] + _, run, _, _ = batched( + tmp_path, + "s", + {"a1": "x", "a2": failure, "a3": "y", "b1": "z"}, + workers, + on_batch_done=lambda tasks: seen.append(tasks[0].subject[0]), + ) + assert run.exit_code == 1 + assert seen == ["b"] + + +def test_incomplete_task_descriptions_defer_batch_observation(tmp_path): + class Missing(Batched): + def inputs(self, subject, env): + if subject == "a2": + raise ValueError("unavailable input") + return super().inputs(subject, env) + + log = tmp_path / "log" + log.mkdir() + seen = [] + _, run, _ = orchestrate( + tmp_path, + "s", + {"a1": "x", "a2": "y"}, + stages=registry(Missing(log), Count(), Total()), + on_batch_done=seen.append, + ) + assert run.exit_code == 1 and seen == [] + + +@pytest.mark.parametrize("workers", [0, 2]) +def test_interruption_before_batch_end_does_not_notify(tmp_path, monkeypatch, workers): + depot = Depot(tmp_path / "s") + log = tmp_path / "log" + log.mkdir() + seen = [] + orch = Orchestrator(depot, registry(Batched(log)), workers=workers, on_batch_done=seen.append, log=lambda _: None) + done = orch._done + + def interrupt(*args): + done(*args) + raise KeyboardInterrupt + + monkeypatch.setattr(orch, "_done", interrupt) + with pytest.raises(KeyboardInterrupt): + orch.run(["toy.src"], facts={"files": toy_files(tmp_path / "in", {"a1": "x", "a2": "y"})}) + assert seen == [] and orch._pool == [] + + +def test_configuration_stop_keeps_the_incomplete_batch(tmp_path): + from nnnotes.source.config import ConfigError + + class Stops(Batched): + def run(self, task, depot): + if task.subject == "a2": + raise ConfigError("setting paths.tool is not set") + return super().run(task, depot) + + log = tmp_path / "log" + log.mkdir() + seen = [] + orch = Orchestrator(Depot(tmp_path / "s"), registry(Stops(log)), on_batch_done=seen.append, log=lambda _: None) + with pytest.raises(ConfigError): + orch.run(["toy.src"], facts={"files": toy_files(tmp_path / "in", {"a1": "x", "a2": "y", "a3": "z"})}) + assert seen == [] + + def test_jobs_group_a_stages_batches(): def task(stage, name, cpu, peak): return Task(stage, 1, name, cost=Cost(cpu, peak))