diff --git a/TESTING.md b/TESTING.md index 7a7e9630..60491a60 100644 --- a/TESTING.md +++ b/TESTING.md @@ -3900,12 +3900,11 @@ fallback (`_get_fallback_prompt` is gone) and no unproven tracker entry. `positive_prompt_id` hint). The monitor tries it last; the scan tries it when the #170 browser-parser floor finds nothing, and also for files with no prompt metadata at all (`--disable-metadata`, savers that embed none); it never - creates a prompt. The lookup uses `idx_history_read_execution_run_outputs` - (`filename, subfolder, root_key, path_style`), added by - `migrate_execution_history`; the output schema check accepts exactly that - optional index and still refuses any other. An older build refuses a library - that has it, so downgrading after this change needs `DROP INDEX - idx_history_read_execution_run_outputs` first. + creates a prompt. Lookups are batched (`stored_run_prompts`): one + parameterized `filename IN (...)` query per chunk of at most 500 names + (`LOOKUP_CHUNK`), matched in memory, so a scan batch asks history once per + chunk, never once per file. The history schema is unchanged (no extra index), + so older builds still open the library. - **After the run ends.** An image with no run keys (no embedded API graph) that nothing identifies waits for its run: when the run's outputs are recorded (`write_execution_outputs_for_finish` notifies `subscribe_outputs_sealed` diff --git a/database/execution_history_schema.py b/database/execution_history_schema.py index 0f04f04b..191033d3 100644 --- a/database/execution_history_schema.py +++ b/database/execution_history_schema.py @@ -314,16 +314,12 @@ def migrate_execution_history(conn: sqlite3.Connection) -> None: if version == 1: install_context_schema(conn) install_output_schema(conn) - from .execution_output_schema import OUTPUT_READ_INDEX_COLUMNS - # Read projections need reverse content lookup and the complete run order. for table, columns in ( ("execution_prompt_snapshots", "raw_text_id,snapshot_id"), ("execution_contexts", "positive_snapshot_id,context_id"), ("execution_run_contexts", "context_id,run_id,anchor_id"), ("execution_runs", "observed_dequeue_at,session_id,queue_item_id,run_id"), - # Attribution of an output file to its run (#342) - ("execution_run_outputs", ",".join(OUTPUT_READ_INDEX_COLUMNS)), ): conn.execute( f"CREATE INDEX IF NOT EXISTS idx_history_read_{table} " diff --git a/database/execution_output_schema.py b/database/execution_output_schema.py index b6e06381..c39867e5 100644 --- a/database/execution_output_schema.py +++ b/database/execution_output_schema.py @@ -119,25 +119,6 @@ def install_output_schema(conn: sqlite3.Connection) -> None: conn.execute(_table_sql(table)) -# The one optional, non-unique read index (#342): output file -> run. Older -# libraries may lack it; migrate_execution_history adds it. -OUTPUT_READ_INDEX = "idx_history_read_execution_run_outputs" -OUTPUT_READ_INDEX_COLUMNS = ("filename", "subfolder", "root_key", "path_style") - - -def _is_read_index(table: str, index: tuple, keys: tuple) -> bool: - """Exactly the attribution read index, as the migration creates it.""" - return ( - table == "execution_run_outputs" - and index[1] == OUTPUT_READ_INDEX - and not index[2] - and index[3] == "c" - and not index[4] - and tuple(key[2] for key in keys) == OUTPUT_READ_INDEX_COLUMNS - and not any(key[3] or key[4] != "BINARY" for key in keys) - ) - - def output_schema_error(conn: sqlite3.Connection) -> str | None: """Reject altered schemas and invalid rows before returning bounded DTOs.""" if conn.execute( @@ -176,8 +157,6 @@ def output_schema_error(conn: sqlite3.Connection) -> str | None: keys = tuple( row for row in conn.execute(f'PRAGMA index_xinfo("{name}")') if row[5] ) - if _is_read_index(table, index, keys): - continue if ( not index[2] or index[3] != "pk" diff --git a/database/run_attribution.py b/database/run_attribution.py index 97f19f35..8917a6ec 100644 --- a/database/run_attribution.py +++ b/database/run_attribution.py @@ -14,6 +14,11 @@ only when all of this holds: - every positive context of that run names the same stored prompt. Anything else returns None: the caller leaves the file unlinked. Read-only. + +Lookups are batched (``stored_run_prompts``): one parameterized query per +chunk of at most ``LOOKUP_CHUNK`` distinct file names, matched in memory. The +history schema is left exactly as it is (no extra index), so older builds can +still open the library. """ import logging @@ -21,7 +26,7 @@ import os import sqlite3 import threading import weakref -from typing import Any, Callable, Iterable, List, Optional +from typing import Any, Callable, Dict, Iterable, List, Optional try: from ..utils.execution_outputs import snapshot_output_roots @@ -34,12 +39,15 @@ except ImportError: # system clocks, the write finishing just after the host reported it). MTIME_SLACK_MS = 2000 -_RUNS_FOR_OUTPUT = ( - "SELECT DISTINCT o.run_id, r.observed_dequeue_at, " - "COALESCE(r.terminal_event_at, r.observed_callback_at) " +# Distinct file names per lookup query: well under SQLite's bound-parameter +# limit (999 on old builds), so a scan batch costs one query per chunk. +LOOKUP_CHUNK = 500 + +_RUNS_FOR_NAMES = ( + "SELECT DISTINCT o.filename, o.subfolder, o.root_key, o.path_style, o.run_id, " + "r.observed_dequeue_at, COALESCE(r.terminal_event_at, r.observed_callback_at) " "FROM execution_run_outputs o JOIN execution_runs r ON r.run_id = o.run_id " - "WHERE o.folder_type = 'output' AND o.path_style = ? AND o.root_key = ? " - "AND o.subfolder = ? AND o.filename = ? LIMIT 2" + "WHERE o.folder_type = 'output' AND o.filename IN ({marks})" ) @@ -93,26 +101,32 @@ def stored_run_prompt( path_style: Optional[str] = None, ) -> Optional[int]: """The stored prompt id history attributes to ``image_path``, else None.""" - if not image_path or not output_root: - return None - style = path_style or ("windows" if os.name == "nt" else "posix") - location = _location(image_path, output_root, style) - if location is None: - return None - root_key, subfolder, filename = location - try: - rows = ( - db.model.get_connection() - .execute(_RUNS_FOR_OUTPUT, (style, root_key, subfolder, filename)) - .fetchall() - ) - except sqlite3.Error: - return None # no history tables (a library without history) - if len(rows) != 1: - return None - run_id, started, finished = rows[0][0], rows[0][1], rows[0][2] - if not _written_during(image_path, started, finished): + if not image_path: return None + return stored_run_prompts(db, [image_path], output_root, path_style).get(image_path) + + +def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]: + """``{(filename, subfolder, root_key, style): {run_id: (start, end)}}``. + + One query per ``LOOKUP_CHUNK`` names; None without history tables. + """ + found: Dict[tuple, dict] = {} + conn = db.model.get_connection() + for start in range(0, len(names), LOOKUP_CHUNK): + chunk = names[start : start + LOOKUP_CHUNK] + query = _RUNS_FOR_NAMES.format(marks=",".join("?" for _ in chunk)) + try: + rows = conn.execute(query, chunk).fetchall() # nosec B608 + except sqlite3.Error: + return None # no history tables (a library without history) + for name, subfolder, root_key, style, run_id, began, ended in rows: + runs = found.setdefault((name, subfolder, root_key, style), {}) + runs[run_id] = (began, ended) + return found + + +def _named_prompt(db: Any, run_id: str) -> Optional[int]: contexts = db.get_execution_contexts(run_id).get("contexts") or [] named = { context.get("positive_prompt_id") @@ -125,6 +139,47 @@ def stored_run_prompt( return prompt_id if isinstance(prompt_id, int) else None +def stored_run_prompts( + db: Any, + image_paths: Iterable[str], + output_root: Optional[str], + path_style: Optional[str] = None, +) -> Dict[str, Optional[int]]: + """``{path: stored prompt id or None}`` for many files in few queries. + + The output lookup is one query per ``LOOKUP_CHUNK`` distinct names; each + run that alone reported a file is then read once for its contexts. + """ + paths = list(dict.fromkeys(path for path in image_paths if path)) + result: Dict[str, Optional[int]] = dict.fromkeys(paths) + if not paths or not output_root: + return result + style = path_style or ("windows" if os.name == "nt" else "posix") + located = {} + for path in paths: + location = _location(path, output_root, style) + if location is not None: + located[path] = location + if not located: + return result + names = sorted({filename for _key, _sub, filename in located.values()}) + runs = _runs_by_location(db, names) + if not runs: + return result + prompts: Dict[str, Optional[int]] = {} + for path, (root_key, subfolder, filename) in located.items(): + reported = runs.get((filename, subfolder, root_key, style), {}) + if len(reported) != 1: + continue + ((run_id, (began, ended)),) = reported.items() + if not _written_during(path, began, ended): + continue + if run_id not in prompts: + prompts[run_id] = _named_prompt(db, run_id) + result[path] = prompts[run_id] + return result + + _LOGGER = logging.getLogger("prompt_manager.run_attribution") _SEALED_LISTENERS: List[weakref.WeakMethod] = [] _SEALED_LOCK = threading.Lock() diff --git a/py/api/admin.py b/py/api/admin.py index d9fc7b0b..2c307542 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -36,7 +36,11 @@ except ImportError: try: from ...database.pair_ingestion import blank_positive, ingest_resolved_pair - from ...database.run_attribution import host_output_directory, stored_run_prompt + from ...database.run_attribution import ( + host_output_directory, + stored_run_prompt, + stored_run_prompts, + ) from ...utils.generation_params import ( PLACEHOLDERS, GenerationParamsError, @@ -47,7 +51,11 @@ try: from ...utils.review_capture import reviewable except ImportError: from database.pair_ingestion import blank_positive, ingest_resolved_pair - from database.run_attribution import host_output_directory, stored_run_prompt + from database.run_attribution import ( + host_output_directory, + stored_run_prompt, + stored_run_prompts, + ) from utils.generation_params import ( PLACEHOLDERS, GenerationParamsError, @@ -1566,6 +1574,8 @@ class AdminRoutesMixin: if outcomes is not None: outcomes.append(_scan_outcome(media_file, roots, status, title)) + batch_results = list(batch_results) + history = self._recorded_runs_for(media for media, _meta in batch_results) for media_file, metadata in batch_results: if self._scan_job and self._scan_job.cancelled: break @@ -1579,6 +1589,7 @@ class AdminRoutesMixin: ambiguous=ambiguous, untraced=untraced, records=records, + history=history, ) except Exception as e: self.logger.error(f"Error processing {media_file.name}: {e}") @@ -1637,6 +1648,7 @@ class AdminRoutesMixin: ambiguous=None, untraced=None, records=None, + history=None, ): """Enrich exact pairs; retain legacy positive-only import for absent pairs. @@ -1656,7 +1668,7 @@ class AdminRoutesMixin: if not metadata: # No prompt metadata (--disable-metadata, bare savers): only the # recorded run can say which prompt made it (#342) - return self._link_recorded_run(media_file, {}, detailed) + return self._link_recorded_run(media_file, {}, detailed, history) canonical = dict(metadata) for key in ("prompt", "workflow"): if isinstance(canonical.get(key), str): @@ -1666,7 +1678,7 @@ class AdminRoutesMixin: # Preserve malformed metadata for conservative resolver rejection. pass if not canonical.get("prompt") and not canonical.get("parameters"): - return self._link_recorded_run(media_file, canonical, detailed) + return self._link_recorded_run(media_file, canonical, detailed, history) _with_file_size(canonical, media_file) resolution = resolve_prompt_pair(canonical) if blank_positive(resolution): @@ -1700,7 +1712,9 @@ class AdminRoutesMixin: # Graphs without a recognized sampler keep the legacy extractor below. prompt_text = self._graph_prompt_text(metadata) if not prompt_text: - linked = self._link_recorded_run(media_file, canonical, detailed) + linked = self._link_recorded_run( + media_file, canonical, detailed, history + ) if linked is not None: return linked return "unresolved" if unchoosable else resolution["state"] @@ -1726,23 +1740,42 @@ class AdminRoutesMixin: return status return {"status": status, "title": self._scan_title(media_file, prompt_text)} - def _link_recorded_run(self, media_file, canonical, detailed=False): + def _recorded_runs_for(self, media_files): + """``{path: prompt id or None}`` from recorded history for a scan batch. + + One output lookup per chunk of names (``stored_run_prompts``), so a + batch never queries history once per file. An empty dict on failure. + """ + try: + return stored_run_prompts( + self.db, [str(path) for path in media_files], host_output_directory() + ) + except Exception: + self.logger.warning("History attribution failed for a scan batch") + return {} + + def _link_recorded_run(self, media_file, canonical, detailed=False, history=None): """Link a file recorded history attributes to one stored prompt (#342). Returns the scan status (with ``detailed``, a stored or linked file's ``{"status", "title"}``), or None when history names no single stored prompt for the file (``database.run_attribution``). Never creates a - prompt. Works without any prompt metadata in the file. + prompt. Works without any prompt metadata in the file. ``history`` is + the batch's precomputed lookup; without it the file is looked up alone. """ - try: - prompt_id = stored_run_prompt( - self.db, str(media_file), host_output_directory() - ) - except Exception: - self.logger.warning( - f"History attribution failed for {media_file.name}", exc_info=True - ) - return None + if history is not None: + prompt_id = history.get(str(media_file)) + else: + try: + prompt_id = stored_run_prompt( + self.db, str(media_file), host_output_directory() + ) + except Exception: + self.logger.warning( + f"History attribution failed for {media_file.name}", + exc_info=True, + ) + return None if prompt_id is None: return None try: diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index a2f8e586..a5b85ba0 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -2466,6 +2466,7 @@ class TestActivityJobs(AdminAPITestCase): ambiguous=None, untraced=None, records=None, + history=None, ): key = Path(media_file).relative_to(self.output_dir).as_posix() outcome = names[key] diff --git a/tests/test_run_attribution.py b/tests/test_run_attribution.py index b59c6959..5ba9e695 100644 --- a/tests/test_run_attribution.py +++ b/tests/test_run_attribution.py @@ -440,53 +440,99 @@ class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase): run_attribution.notify_outputs_sealed("run", []) # no dead callback raises -class TestOutputLookupIndex(unittest.TestCase): - """M5: the history lookup by file name uses an index, added by migration.""" +class TestBoundedLookup(RecordedRuns, unittest.TestCase): + """M5: history lookups are batched, and the history schema is unchanged.""" setUp = history_tests.HistoryDatabaseTests.setUp close_databases = history_tests.HistoryDatabaseTests.close_databases database = history_tests.HistoryDatabaseTests.database - def test_index_exists_is_used_validated_and_reinstalled(self): - from database.execution_history_schema import ( - history_schema_error, - migrate_execution_history, - ) - from database.run_attribution import _RUNS_FOR_OUTPUT + def output_lookups(self, db, action): + statements = [] + conn = db.model.get_connection() + conn.set_trace_callback(statements.append) + try: + result = action() + finally: + conn.set_trace_callback(None) + return result, [s for s in statements if "FROM execution_run_outputs" in s] + + def test_one_query_per_chunk_not_per_file_and_same_answers(self): + from database import run_attribution db = self.database() - conn = db.model.get_connection() - name = "idx_history_read_execution_run_outputs" - columns = [r[2] for r in conn.execute(f"PRAGMA index_info({name})")] - self.assertEqual(columns, ["filename", "subfolder", "root_key", "path_style"]) - plan = " ".join( - str(row[3]) - for row in conn.execute( - "EXPLAIN QUERY PLAN " + _RUNS_FOR_OUTPUT, ("posix", "k", "", "f.png") + output = str(Path(self.temp.name) / "output") + os.makedirs(output) + prompt_id = db.save_prompt( + text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT) + ) + names = [f"img_{n:05}_.png" for n in range(25)] + paths = [os.path.join(output, name) for name in names] + for path in paths: + Path(path).write_bytes(b"png") + self.record_run(db, output, [("", n) for n in names[:10]], [LLM_TEXT]) + self.record_run(db, output, [("", n) for n in names[5:10]], [""], item=1) + with mock.patch.object(run_attribution, "LOOKUP_CHUNK", 10): + result, lookups = self.output_lookups( + db, lambda: run_attribution.stored_run_prompts(db, paths, output) + ) + self.assertEqual(len(lookups), 3) # 25 names in chunks of 10 + expected = {path: None for path in paths} + expected.update({path: prompt_id for path in paths[:5]}) + self.assertEqual(result, expected) + for path in paths[:12]: + self.assertEqual( + run_attribution.stored_run_prompt(db, path, output), expected[path] ) - ) - self.assertIn(name, plan) - self.assertIsNone(history_schema_error(conn, allow_missing=False)) - conn.execute(f"DROP INDEX {name}") - self.assertIsNone(history_schema_error(conn, allow_missing=False)) - migrate_execution_history(conn) - migrate_execution_history(conn) - self.assertTrue( - conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (name,)).fetchone() - ) - def test_any_other_index_on_the_outputs_table_is_still_refused(self): - from database.execution_output_schema import output_schema_error + def test_the_history_schema_gains_no_index(self): + from database.execution_history_schema import history_schema_error db = self.database() conn = db.model.get_connection() - conn.execute("CREATE INDEX idx_rogue ON execution_run_outputs(filename)") - self.assertIn("unexpected", output_schema_error(conn)) - conn.execute("DROP INDEX idx_rogue") - name = "idx_history_read_execution_run_outputs" - conn.execute(f"DROP INDEX {name}") - conn.execute(f"CREATE INDEX {name} ON execution_run_outputs(filename)") - self.assertIn("unexpected", output_schema_error(conn)) + indexes = [ + row[1] + for row in conn.execute("PRAGMA index_list(execution_run_outputs)") + if row[3] != "pk" + ] + self.assertEqual(indexes, []) + self.assertIsNone(history_schema_error(conn, allow_missing=False)) + + +class TestScanBatchLookup(RecordedRuns, AdminAPITestCase): + """M5: a scan batch asks history once per chunk of names, not per file.""" + + def test_metadata_less_batch_links_with_one_history_lookup(self): + db = self.api.db + prompt_id = db.save_prompt( + text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT) + ) + files = [] + for n in range(6): + path = self.output_dir / f"bare_{n:05}_.png" + path.write_bytes(b"png") + files.append(path) + self.record_run( + db, str(self.output_dir), [("", f.name) for f in files[:4]], [LLM_TEXT] + ) + statements = [] + conn = db.model.get_connection() + conn.set_trace_callback(statements.append) + try: + counts = self.api._ingest_scan_batch_sync([(f, {}) for f in files]) + finally: + conn.set_trace_callback(None) + lookups = [s for s in statements if "FROM execution_run_outputs" in s] + self.assertEqual(len(lookups), 1) + self.assertEqual(counts["linked"], 4) + linked = { + row[0] + for row in conn.execute( + "SELECT filename FROM generated_images WHERE prompt_id=?", + (prompt_id,), + ) + } + self.assertEqual(linked, {f.name for f in files[:4]}) if __name__ == "__main__":