diff --git a/Dockerfile.multistage b/Dockerfile.multistage index 1d163826..da2be129 100644 --- a/Dockerfile.multistage +++ b/Dockerfile.multistage @@ -199,7 +199,7 @@ RUN echo "[+] Initializing image collection..." \ RUN chmod +x "$CODE_DIR"/bin/*.sh \ && chown -R "$DEFAULT_PUID:$DEFAULT_PGID" "$LIB_DIR" \ && chmod g+w "$TMP_DIR" "$LIB_DIR" "$PLAYWRIGHT_BROWSERS_PATH" \ - && GIT_BINARY="$LIB_DIR/env/bin/git" GALLERYDL_BINARY="$LIB_DIR/env/bin/gallery-dl" FORUMDL_BINARY="$LIB_DIR/env/bin/forum-dl" ABXPKG_INSTALL_TIMEOUT=600 ABXPKG_POSTINSTALL_SCRIPTS=True ABXPKG_MIN_RELEASE_AGE=0 TIMEOUT=600 gosu "$ARCHIVEBOX_USER" archivebox install archivewebpage defuddle forumdl gallerydl git istilldontcareaboutcookies liteparse mercury papersdl parse_rss_urls readability search_backend_ripgrep search_backend_sonic 2>&1 | tee -a /VERSION.txt \ + && GIT_BINARY="$LIB_DIR/env/bin/git" GALLERYDL_BINARY="$LIB_DIR/env/bin/gallery-dl" FORUMDL_BINARY="$LIB_DIR/env/bin/forum-dl" ABXPKG_INSTALL_TIMEOUT=600 ABXPKG_POSTINSTALL_SCRIPTS=True ABXPKG_MIN_RELEASE_AGE=0 TIMEOUT=600 gosu "$ARCHIVEBOX_USER" archivebox install archivewebpage defuddle forumdl gallerydl git istilldontcareaboutcookies liteparse mercury opendataloader papersdl parse_rss_urls readability search_backend_ripgrep search_backend_sonic 2>&1 | tee -a /VERSION.txt \ && "$LIB_DIR/env/bin/chromium" --version | tee -a /VERSION.txt \ && "$LIB_DIR/pip/packages/papers-dl/venv/bin/papers-dl" --version | tee -a /VERSION.txt \ && /usr/bin/rg --version | head -1 | tee -a /VERSION.txt \ diff --git a/archivebox/cli/archivebox_extract.py b/archivebox/cli/archivebox_extract.py index c315401e..cbd950d4 100644 --- a/archivebox/cli/archivebox_extract.py +++ b/archivebox/cli/archivebox_extract.py @@ -116,6 +116,7 @@ def run_plugins( wait: bool = True, emit_results: bool = True, show_progress: bool = True, + preserve_queued: bool = False, ) -> int: """ Run plugins on Snapshots from input. @@ -219,6 +220,18 @@ def run_plugins( else: requested_rows.add((snapshot_id, plugin_name, "")) + queued_rows: set[tuple[str, str, str]] = set() + if preserve_queued and requested_rows: + queued_rows = { + (str(snapshot_id), plugin_name, hook_name) + for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( + snapshot_id__in=existing_snapshot_ids, + plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in requested_rows}, + status=ArchiveResult.StatusChoices.QUEUED, + ).values_list("snapshot_id", "plugin", "hook_name") + } + rows_to_queue = requested_rows - queued_rows + reset_fields = { "status": ArchiveResult.StatusChoices.QUEUED, "output_str": "", @@ -230,23 +243,34 @@ def run_plugins( "end_ts": None, "modified_at": timezone.now(), } - if plugins_list: - ArchiveResult.objects.filter(snapshot_id__in=existing_snapshot_ids, plugin__in=plugins_list).update(**reset_fields) - elif requested_plugins_by_snapshot: - snapshot_ids_by_plugin: dict[str, set[str]] = defaultdict(set) - for snapshot_id, plugin_names in requested_plugins_by_snapshot.items(): - if snapshot_id in existing_snapshot_ids: - for plugin_name in plugin_names: - snapshot_ids_by_plugin[plugin_name].add(snapshot_id) - for plugin_name, plugin_snapshot_ids in snapshot_ids_by_plugin.items(): - ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name).update(**reset_fields) - existing_rows = set( - ArchiveResult.objects.filter( - snapshot_id__in=existing_snapshot_ids, - plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in requested_rows}, - ).values_list("snapshot_id", "plugin", "hook_name"), + if rows_to_queue and plugins_list: + rows_to_reset_by_hook: dict[tuple[str, str], set[str]] = defaultdict(set) + for snapshot_id, plugin_name, hook_name in rows_to_queue: + rows_to_reset_by_hook[(plugin_name, hook_name)].add(snapshot_id) + for (plugin_name, hook_name), plugin_snapshot_ids in rows_to_reset_by_hook.items(): + ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name, hook_name=hook_name).update( + **reset_fields, + ) + elif rows_to_queue and requested_plugins_by_snapshot: + snapshot_ids_by_hook: dict[tuple[str, str], set[str]] = defaultdict(set) + for snapshot_id, plugin_name, hook_name in rows_to_queue: + snapshot_ids_by_hook[(plugin_name, hook_name)].add(snapshot_id) + for (plugin_name, hook_name), plugin_snapshot_ids in snapshot_ids_by_hook.items(): + ArchiveResult.objects.filter(snapshot_id__in=plugin_snapshot_ids, plugin=plugin_name, hook_name=hook_name).update( + **reset_fields, + ) + existing_rows = ( + { + (str(snapshot_id), plugin_name, hook_name) + for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( + snapshot_id__in=existing_snapshot_ids, + plugin__in={plugin_name for _snapshot_id, plugin_name, _hook_name in rows_to_queue}, + ).values_list("snapshot_id", "plugin", "hook_name") + } + if rows_to_queue + else set() ) - missing_rows = requested_rows - {(str(snapshot_id), plugin_name, hook_name) for snapshot_id, plugin_name, hook_name in existing_rows} + missing_rows = rows_to_queue - existing_rows if missing_rows: ArchiveResult.objects.bulk_create( [ @@ -271,7 +295,23 @@ def run_plugins( # also keep status=paused here: `retry_at` only asks the orchestrator # to process the queued plugin rows, and run_due_snapshot restores # retry_at=MAX afterward instead of resuming the snapshot lifecycle. - for snapshot in Snapshot.objects.filter(id__in=existing_snapshot_ids).only("id", "status", "modified_at"): + affected_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in rows_to_queue} + if preserve_queued and queued_rows: + queued_snapshot_ids = {snapshot_id for snapshot_id, _plugin_name, _hook_name in queued_rows} + affected_snapshot_ids.update( + str(snapshot_id) + for snapshot_id in Snapshot.objects.filter(id__in=queued_snapshot_ids) + .filter(retry_at__gt=queue_at) + .values_list("id", flat=True) + ) + affected_snapshot_ids.update( + str(snapshot_id) + for snapshot_id in Snapshot.objects.filter(id__in=queued_snapshot_ids, retry_at__isnull=True).values_list( + "id", + flat=True, + ) + ) + for snapshot in Snapshot.objects.filter(id__in=affected_snapshot_ids).only("id", "status", "modified_at"): # Guard the read-time status so we never bump retry_at on a # row that's been re-queued / started by a concurrent runner. snapshot.safe_update( diff --git a/archivebox/cli/archivebox_update.py b/archivebox/cli/archivebox_update.py index 83741116..3adb935c 100644 --- a/archivebox/cli/archivebox_update.py +++ b/archivebox/cli/archivebox_update.py @@ -105,9 +105,17 @@ def reindex_snapshots( wait_for_turn=None, ) -> dict[str, Any]: from archivebox.cli.archivebox_extract import run_plugins + from archivebox.core.models import ArchiveResult + from abx_dl.models import discover_plugins - stats: dict[str, Any] = {"processed": 0, "queued": 0, "reindexed": 0, "snapshot_ids": []} + stats: dict[str, Any] = {"processed": 0, "requested": 0, "queued": 0, "skipped_queued": 0, "reindexed": 0, "snapshot_ids": []} records: list[dict[str, str]] = [] + plugins_by_name = discover_plugins() + required_hooks_by_plugin = { + plugin_name: frozenset(hook.name for hook in plugins_by_name[plugin_name].filter_hooks("Snapshot")) + for plugin_name in search_plugins + if plugin_name in plugins_by_name + } total = snapshots.count() print(f"[*] Reindexing {total} snapshots with search plugins: {', '.join(search_plugins)}") @@ -118,6 +126,31 @@ def reindex_snapshots( if wait_for_turn: wait_for_turn() batch_records = list(records) + snapshot_ids = {record["snapshot_id"] for record in batch_records} + plugin_names = {record["plugin"] for record in batch_records} + queued_rows = { + (str(snapshot_id), plugin_name, hook_name) + for snapshot_id, plugin_name, hook_name in ArchiveResult.objects.filter( + snapshot_id__in=snapshot_ids, + plugin__in=plugin_names, + status=ArchiveResult.StatusChoices.QUEUED, + ).values_list("snapshot_id", "plugin", "hook_name") + } + records_to_queue = [] + for record in batch_records: + snapshot_id = record["snapshot_id"] + plugin_name = record["plugin"] + required_hooks = required_hooks_by_plugin.get(plugin_name, frozenset()) + if required_hooks and all((snapshot_id, plugin_name, hook_name) in queued_rows for hook_name in required_hooks): + stats["skipped_queued"] += 1 + continue + records_to_queue.append(record) + if not records_to_queue: + print( + f" [{stats['processed']}/{total}] Already queued {len(batch_records)} index jobs", + ) + records.clear() + return # `archivebox update --index-only` intentionally breaks the usual # "runner discovers work" rule by inserting synthetic queued # ArchiveResult rows for search backends. run_plugins() keeps this as @@ -127,15 +160,17 @@ def reindex_snapshots( # finish. exit_code = run_plugins( args=(), - records=batch_records, + records=records_to_queue, wait=False, emit_results=False, show_progress=False, + preserve_queued=True, ) if exit_code != 0: raise SystemExit(exit_code) + stats["queued"] += len(records_to_queue) print( - f" [{stats['processed']}/{total}] Queued {len(batch_records)} index jobs for orchestrator", + f" [{stats['processed']}/{total}] Queued {len(records_to_queue)} index jobs for orchestrator", ) records.clear() @@ -156,7 +191,7 @@ def reindex_snapshots( "plugin": plugin_name, }, ) - stats["queued"] += 1 + stats["requested"] += 1 if len(records) >= batch_size: run_batch() except KeyboardInterrupt as err: @@ -366,15 +401,55 @@ def update( after=after, resume=resume, ) - stats = reindex_snapshots( - snapshots, - search_plugins=search_plugins, - batch_size=batch_size, - collect_ids=is_filtered_update, - wait_for_turn=wait_for_turn, + from django.db.models import Exists, OuterRef, Q + from django.utils import timezone + from archivebox.core.models import ArchiveResult, Snapshot + + scoped_snapshot_ids = snapshots.order_by().values("id") if is_filtered_update else None + queued_index_results = ArchiveResult.objects.filter( + status=ArchiveResult.StatusChoices.QUEUED, + plugin__in=search_plugins, ) - print_index_stats(stats) - touched_snapshot_ids.update(stats.get("snapshot_ids", [])) + if scoped_snapshot_ids is not None: + queued_index_results = queued_index_results.filter(snapshot_id__in=scoped_snapshot_ids) + + if queued_index_results.exists(): + now = timezone.now() + queued_result_for_snapshot = queued_index_results.filter(snapshot_id=OuterRef("pk")) + snapshots_to_wake = ( + Snapshot.objects.filter( + status__in=(Snapshot.StatusChoices.SEALED, Snapshot.StatusChoices.PAUSED), + ) + .annotate( + has_queued_index_result=Exists(queued_result_for_snapshot), + ) + .filter( + has_queued_index_result=True, + ) + .filter( + Q(retry_at__isnull=True) | Q(retry_at__gt=now), + ) + ) + if scoped_snapshot_ids is not None: + snapshots_to_wake = snapshots_to_wake.filter(id__in=scoped_snapshot_ids) + woken_count = snapshots_to_wake.update( + retry_at=now, + modified_at=now, + ) + print( + "[*] Existing queued search index jobs found; " + f"skipping backfill scan and waking {woken_count} snapshot(s) for the runner.", + ) + else: + stats = reindex_snapshots( + snapshots, + search_plugins=search_plugins, + batch_size=batch_size, + collect_ids=is_filtered_update, + wait_for_turn=wait_for_turn, + ) + print_index_stats(stats) + touched_snapshot_ids.update(stats.get("snapshot_ids", [])) if do_run_until_idle and (do_index or not ran_post_migrate_runner): # Search/index backfill intentionally queues targeted @@ -904,7 +979,9 @@ def print_index_stats(stats: dict[str, Any]) -> None: print(f""" [green]Search Reindex Complete[/green] Scanned rows: {stats["processed"]} + Requested jobs: {stats.get("requested", stats["queued"])} Queued index jobs: {stats["queued"]} + Already queued: {stats.get("skipped_queued", 0)} """) diff --git a/archivebox/machine/models.py b/archivebox/machine/models.py index 3f683fe2..03babadd 100755 --- a/archivebox/machine/models.py +++ b/archivebox/machine/models.py @@ -548,6 +548,7 @@ class Binary(ModelWithHealthStats, ModelWithStateMachine): state_machine_name: str | None = "archivebox.machine.models.BinaryMachine" active_state: str = StatusChoices.QUEUED + warn_on_save_outside_runner = False objects = BinaryManager() # pyright: ignore[reportIncompatibleVariableOverride] diff --git a/archivebox/workers/models.py b/archivebox/workers/models.py index aca751e0..771d0331 100644 --- a/archivebox/workers/models.py +++ b/archivebox/workers/models.py @@ -65,6 +65,7 @@ class BaseModelWithStateMachine(models.Model): state_field_name: str state_machine_attr: str = "sm" bind_events_as_methods: bool = False + warn_on_save_outside_runner: ClassVar[bool] = True active_state: ObjectState retry_at_field_name: str @@ -297,7 +298,7 @@ class BaseModelWithStateMachine(models.Model): from archivebox.machine.models import Process process = Process.current() - if process.process_type != Process.TypeChoices.ORCHESTRATOR: + if self.warn_on_save_outside_runner and process.process_type != Process.TypeChoices.ORCHESTRATOR: root_type = getattr(process.root, "process_type", None) if root_type != Process.TypeChoices.ORCHESTRATOR: caller = ""