From 1f0f86e3a2d79cce9352ff28e18a85be2cb625f5 Mon Sep 17 00:00:00 2001 From: Karim shoair Date: Mon, 19 Jan 2026 03:43:10 +0200 Subject: [PATCH] feat(spiders): Add pause/resume system for crawls --- scrapling/spiders/checkpoint.py | 90 +++++++++++++++++++++++++++ scrapling/spiders/engine.py | 107 +++++++++++++++++++++++++++++--- scrapling/spiders/request.py | 28 ++++++++- scrapling/spiders/result.py | 6 ++ scrapling/spiders/scheduler.py | 36 ++++++++++- scrapling/spiders/spider.py | 97 ++++++++++++++++++++++++----- 6 files changed, 336 insertions(+), 28 deletions(-) create mode 100644 scrapling/spiders/checkpoint.py diff --git a/scrapling/spiders/checkpoint.py b/scrapling/spiders/checkpoint.py new file mode 100644 index 0000000..ae795dc --- /dev/null +++ b/scrapling/spiders/checkpoint.py @@ -0,0 +1,90 @@ +import pickle +from pathlib import Path +from dataclasses import dataclass, field + +import anyio +from anyio import Path as AsyncPath + +from scrapling.core.utils import log +from scrapling.core._types import Set, List, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from scrapling.spiders.request import Request + + +@dataclass +class CheckpointData: + """Container for checkpoint state.""" + + requests: List["Request"] = field(default_factory=list) + seen: Set[str] = field(default_factory=set) + + +class CheckpointManager: + """Manages saving and loading checkpoint state to/from disk.""" + + CHECKPOINT_FILE = "checkpoint.pkl" + + def __init__(self, crawldir: str | Path | AsyncPath, interval: float = 300.0): + self.crawldir = AsyncPath(crawldir) + self._checkpoint_path = self.crawldir / self.CHECKPOINT_FILE + self.interval = interval + if not isinstance(interval, (int, float)): + raise TypeError("Checkpoints interval must be integer or float.") + else: + if interval < 0: + raise ValueError("Checkpoints interval must be greater than 0.") + + async def has_checkpoint(self) -> bool: + """Check if a checkpoint exists.""" + return await self._checkpoint_path.exists() + + async def save(self, data: CheckpointData) -> None: + """Save checkpoint data to disk atomically.""" + await self.crawldir.mkdir(parents=True, exist_ok=True) + + temp_path = self._checkpoint_path.with_suffix(".tmp") + + try: + serialized = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL) + async with await anyio.open_file(temp_path, "wb") as f: + await f.write(serialized) + + await temp_path.rename(self._checkpoint_path) + + log.info(f"Checkpoint saved: {len(data.requests)} requests, {len(data.seen)} seen URLs") + except Exception as e: + # Clean up temp file if it exists + if await temp_path.exists(): + await temp_path.unlink() + log.error(f"Failed to save checkpoint: {e}") + raise + + async def load(self) -> Optional[CheckpointData]: + """Load checkpoint data from disk. + + Returns None if no checkpoint exists or if loading fails. + """ + if not await self.has_checkpoint(): + return None + + try: + async with await anyio.open_file(self._checkpoint_path, "rb") as f: + content = await f.read() + data: CheckpointData = pickle.loads(content) + + log.info(f"Checkpoint loaded: {len(data.requests)} requests, {len(data.seen)} seen URLs") + return data + + except Exception as e: + log.error(f"Failed to load checkpoint (starting fresh): {e}") + return None + + async def cleanup(self) -> None: + """Delete checkpoint file after successful completion.""" + try: + if await self._checkpoint_path.exists(): + await self._checkpoint_path.unlink() + log.debug("Checkpoint file cleaned up") + except Exception as e: + log.warning(f"Failed to cleanup checkpoint file: {e}") diff --git a/scrapling/spiders/engine.py b/scrapling/spiders/engine.py index 751d87b..9b226df 100644 --- a/scrapling/spiders/engine.py +++ b/scrapling/spiders/engine.py @@ -1,14 +1,18 @@ import json +import pprint +from pathlib import Path import anyio +from anyio import Path as AsyncPath from anyio import create_task_group, CapacityLimiter, create_memory_object_stream, EndOfStream from scrapling.core.utils import log from scrapling.spiders.request import Request -from scrapling.spiders.result import CrawlStats, ItemList from scrapling.spiders.scheduler import Scheduler from scrapling.spiders.session import SessionManager -from scrapling.core._types import Dict, TYPE_CHECKING, Any, AsyncGenerator +from scrapling.spiders.result import CrawlStats, ItemList +from scrapling.spiders.checkpoint import CheckpointManager, CheckpointData +from scrapling.core._types import Dict, Union, Optional, TYPE_CHECKING, Any, AsyncGenerator if TYPE_CHECKING: from scrapling.spiders.spider import Spider @@ -21,10 +25,16 @@ def _dump(obj: Dict) -> str: class CrawlerEngine: """Orchestrates the crawling process.""" - def __init__(self, spider: "Spider", session_manager: SessionManager, scheduler: Scheduler | None = None): + def __init__( + self, + spider: "Spider", + session_manager: SessionManager, + crawldir: Optional[Union[str, Path, AsyncPath]] = None, + interval: float = 300.0, + ): self.spider = spider self.session_manager = session_manager - self.scheduler = scheduler or Scheduler() + self.scheduler = Scheduler() self.stats = CrawlStats() self._global_limiter = CapacityLimiter(spider.concurrent_requests) @@ -36,6 +46,12 @@ class CrawlerEngine: self._items: ItemList = ItemList() self._item_stream: Any = None + self._checkpoint_system_enabled = bool(crawldir) + self._checkpoint_manager = CheckpointManager(crawldir or "", interval) + self._last_checkpoint_time: float = 0.0 + self._pause_requested: bool = False + self.paused: bool = False + def _is_domain_allowed(self, request: Request) -> bool: """Check if the request's domain is in allowed_domains.""" if not self._allowed_domains: @@ -105,7 +121,7 @@ class CrawlerEngine: processed_result = await self.spider.on_scraped_item(result) if processed_result: self.stats.items_scraped += 1 - log.debug(f"Scraped from {str(response)}\n{processed_result}") + log.debug(f"Scraped from {str(response)}\n{pprint.pformat(processed_result)}") if self._item_stream: await self._item_stream.send(processed_result) else: @@ -116,6 +132,8 @@ class CrawlerEngine: elif result is not None: log.error(f"Spider must return Request, dict or None, got '{type(result)}' in {request}") except Exception as e: + msg = f"Spider error processing {request}:\n {e}" + log.error(msg, exc_info=e) await self.spider.on_error(request, e) async def _task_wrapper(self, request: Request) -> None: @@ -125,25 +143,95 @@ class CrawlerEngine: finally: self._active_tasks -= 1 + def request_pause(self) -> None: + """Request a graceful pause of the crawl.""" + if not self._pause_requested: + self._pause_requested = True + log.info("Pause requested, waiting for in-flight requests to complete...") + + async def _save_checkpoint(self) -> None: + """Save current state to checkpoint files.""" + requests, seen = self.scheduler.snapshot() + data = CheckpointData(requests=requests, seen=seen) + await self._checkpoint_manager.save(data) + self._last_checkpoint_time = anyio.current_time() + + def _is_checkpoint_time(self) -> bool: + """Check if it's time for the periodic checkpoint.""" + if not self._checkpoint_system_enabled: + return False + + if self._checkpoint_manager.interval == 0: + return False + + current_time = anyio.current_time() + return (current_time - self._last_checkpoint_time) >= self._checkpoint_manager.interval + + async def _restore_from_checkpoint(self) -> bool: + """Attempt to restore state from checkpoint. + + Returns True if successfully restored, False otherwise. + """ + if not self._checkpoint_system_enabled: + raise + + data = await self._checkpoint_manager.load() + if data is None: + return False + + self.scheduler.restore(data) + + # Restore callbacks from spider after scheduler restore + for request in data.requests: + request._restore_callback(self.spider) + + return True + async def crawl(self) -> CrawlStats: """Run the spider and return CrawlStats.""" self._running = True self._items.clear() + self.paused = False + self._pause_requested = False self.stats = CrawlStats(start_time=anyio.current_time()) + # Check for existing checkpoint + resuming = (await self._restore_from_checkpoint()) if self._checkpoint_system_enabled else False + self._last_checkpoint_time = anyio.current_time() + async with self.session_manager: self.stats.concurrent_requests = self.spider.concurrent_requests self.stats.concurrent_requests_per_domain = self.spider.concurrent_requests_per_domain self.stats.download_delay = self.spider.download_delay - await self.spider.on_start() + await self.spider.on_start(resuming=resuming) try: - async for request in self.spider.start_requests(): - await self.scheduler.enqueue(request) + if not resuming: + async for request in self.spider.start_requests(): + await self.scheduler.enqueue(request) + else: + log.info("Resuming from checkpoint, skipping start_requests()") # Process queue async with create_task_group() as tg: while self._running: + # Check for pause request + if self._checkpoint_system_enabled: + if self._pause_requested: + # Wait for active tasks to complete + if self._active_tasks == 0: + await self._save_checkpoint() + self.paused = True + self._running = False + log.info("Spider paused, checkpoint saved") + break + # Wait briefly and check again + await anyio.sleep(0.05) + continue + + if self._is_checkpoint_time(): + await self._save_checkpoint() + if self.scheduler.is_empty: # Empty queue + no active tasks = done if self._active_tasks == 0: @@ -161,6 +249,9 @@ class CrawlerEngine: finally: await self.spider.on_close() + # Clean up checkpoint files on successful completion (not paused) + if not self.paused and self._checkpoint_system_enabled: + await self._checkpoint_manager.cleanup() self.stats.log_levels_counter = self.spider._log_counter.get_counts() self.stats.end_time = anyio.current_time() diff --git a/scrapling/spiders/request.py b/scrapling/spiders/request.py index e6b25d0..e2fb588 100644 --- a/scrapling/spiders/request.py +++ b/scrapling/spiders/request.py @@ -1,7 +1,10 @@ from urllib.parse import urlparse from scrapling.engines.toolbelt.custom import Response -from scrapling.core._types import Any, AsyncGenerator, Callable, Dict, Union +from scrapling.core._types import Any, AsyncGenerator, Callable, Dict, Union, TYPE_CHECKING + +if TYPE_CHECKING: + from scrapling.spiders.spider import Spider class Request: @@ -72,3 +75,26 @@ class Request: if not isinstance(other, Request): return NotImplemented return self._fp == other._fp + + def __getstate__(self) -> dict[str, Any]: + """Prepare state for pickling - store callback as name string for pickle compatibility.""" + state = self.__dict__.copy() + state["_callback_name"] = getattr(self.callback, "__name__", None) if self.callback is not None else None + state["callback"] = None # Don't pickle the actual callable + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore state from pickle - callback restored later via _restore_callback().""" + self._callback_name: str | None = state.pop("_callback_name", None) + self.__dict__.update(state) + + def _restore_callback(self, spider: "Spider") -> None: + """Restore callback from spider after unpickling. + + :param spider: Spider instance to look up callback method on + """ + if hasattr(self, "_callback_name") and self._callback_name: + self.callback = getattr(spider, self._callback_name, None) or spider.parse + del self._callback_name + elif hasattr(self, "_callback_name"): + del self._callback_name diff --git a/scrapling/spiders/result.py b/scrapling/spiders/result.py index 7e38a81..eeb5522 100644 --- a/scrapling/spiders/result.py +++ b/scrapling/spiders/result.py @@ -112,6 +112,12 @@ class CrawlResult: stats: CrawlStats items: ItemList + paused: bool = False + + @property + def completed(self) -> bool: + """True if the crawl completed normally (not paused).""" + return not self.paused def __len__(self) -> int: return len(self.items) diff --git a/scrapling/spiders/scheduler.py b/scrapling/spiders/scheduler.py index 2890e50..e36472e 100644 --- a/scrapling/spiders/scheduler.py +++ b/scrapling/spiders/scheduler.py @@ -3,6 +3,10 @@ from itertools import count from scrapling.core.utils import log from scrapling.spiders.request import Request +from scrapling.core._types import List, Set, Tuple, TYPE_CHECKING + +if TYPE_CHECKING: + from scrapling.spiders.checkpoint import CheckpointData class Scheduler: @@ -17,6 +21,8 @@ class Scheduler: self._queue: asyncio.PriorityQueue[tuple[int, int, Request]] = asyncio.PriorityQueue() self._seen: set[str] = set() self._counter = count() + # Mirror dict for snapshot without draining queue + self._pending: dict[int, tuple[int, int, Request]] = {} async def enqueue(self, request: Request) -> bool: """Add a request to the queue.""" @@ -29,12 +35,16 @@ class Scheduler: self._seen.add(fingerprint) # Negative priority so higher priority = dequeued first - await self._queue.put((-request.priority, next(self._counter), request)) + counter = next(self._counter) + item = (-request.priority, counter, request) + self._pending[counter] = item + await self._queue.put(item) return True async def dequeue(self) -> Request: """Get the next request to process.""" - _, _, request = await self._queue.get() + _, counter, request = await self._queue.get() + self._pending.pop(counter, None) return request def __len__(self) -> int: @@ -43,3 +53,25 @@ class Scheduler: @property def is_empty(self) -> bool: return self._queue.empty() + + def snapshot(self) -> Tuple[List[Request], Set[str]]: + """Create a snapshot of the current state for checkpoints.""" + sorted_items = sorted(self._pending.values(), key=lambda x: (x[0], x[1])) # Maintain queue order + requests = [item[2] for item in sorted_items] + return requests, self._seen.copy() + + def restore(self, data: "CheckpointData") -> None: + """Restore scheduler state from checkpoint data. + + :param data: CheckpointData containing requests and seen set + """ + self._seen = data.seen.copy() + + # Restore pending requests in order (they're already sorted by priority) + for request in data.requests: + counter = next(self._counter) + item = (-request.priority, counter, request) + self._pending[counter] = item + self._queue.put_nowait(item) + + log.info(f"Scheduler restored: {len(data.requests)} requests, {len(data.seen)} seen") diff --git a/scrapling/spiders/spider.py b/scrapling/spiders/spider.py index aedab1a..bde8f78 100644 --- a/scrapling/spiders/spider.py +++ b/scrapling/spiders/spider.py @@ -1,15 +1,17 @@ +import signal import logging from pathlib import Path from abc import ABC, abstractmethod import anyio +from anyio import Path as AsyncPath from scrapling.spiders.request import Request from scrapling.spiders.engine import CrawlerEngine from scrapling.spiders.session import SessionManager from scrapling.core.utils import set_logger, reset_logger from scrapling.spiders.result import CrawlResult, CrawlStats -from scrapling.core._types import Set, Any, Dict, Optional, TYPE_CHECKING, AsyncGenerator +from scrapling.core._types import Set, Any, Dict, Optional, Union, TYPE_CHECKING, AsyncGenerator BLOCKED_CODES = {401, 403, 407, 429, 444, 500, 502, 503, 504} if TYPE_CHECKING: @@ -82,7 +84,12 @@ class Spider(ABC): logging_date_format: str = "%Y-%m-%d %H:%M:%S" log_file: Optional[str] = None - def __init__(self): + def __init__(self, crawldir: Optional[Union[str, Path, AsyncPath]] = None, interval: float = 300.0): + """Initialize the spider. + + :param crawldir: Directory for checkpoint files. If provided, enables pause/resume. + :param interval: Seconds between periodic checkpoint saves (default 5 minutes). + """ if self.name is None: raise ValueError(f"{self.__class__.__name__} must have a name.") @@ -109,8 +116,12 @@ class Spider(ABC): file_handler.setFormatter(formatter) self.logger.addHandler(file_handler) + self.crawldir: Optional[Path] = Path(crawldir) if crawldir else None + self._interval = interval + self._engine: Optional[CrawlerEngine] = None + self._original_sigint_handler: Any = None + self._session_manager = SessionManager() - self._stream_engine: CrawlerEngine | None = None try: self.configure_sessions(self._session_manager) @@ -143,11 +154,17 @@ class Spider(ABC): async def parse(self, response: "Response") -> AsyncGenerator[Dict[str, Any] | Request | None, None]: """Default callback for processing responses""" raise NotImplementedError(f"{self.__class__.__name__} must implement parse() method") - yield # Make this a generator + yield # Make this a generator for type checkers - async def on_start(self) -> None: - """Called before crawling starts. Override for setup logic.""" - self.logger.debug("Starting spider") + async def on_start(self, resuming: bool = False) -> None: + """Called before crawling starts. Override for setup logic. + + :param resuming: It's enabled if the spider is resuming from a checkpoint, left for the user to use. + """ + if resuming: + self.logger.debug("Resuming spider from checkpoint") + else: + self.logger.debug("Starting spider") async def on_close(self) -> None: """Called after crawling finishes. Override for cleanup logic.""" @@ -159,7 +176,7 @@ class Spider(ABC): Override for custom error handling. """ - self.logger.error(error, exc_info=error) + pass async def on_scraped_item(self, item: Dict[str, Any]) -> Dict[str, Any] | None: """A hook to be overridden by users to do some processing on scraped items, return `None` to drop the item silently.""" @@ -193,13 +210,45 @@ class Spider(ABC): manager.add("default", FetcherSession()) + def pause(self): + """Pause the crawling process if checkpoint system is enabled.""" + if self._engine: + self._engine.request_pause() + else: + raise RuntimeError("Spider doesn't have active crawl to pause, no crawl engine started!") + + def _setup_signal_handler(self) -> None: + """Set up SIGINT handler for graceful pause.""" + + def handler(_signum: int, _frame: Any) -> None: + if self._engine: + self._engine.request_pause() + else: + # No engine yet, just raise KeyboardInterrupt + raise KeyboardInterrupt + + try: + self._original_sigint_handler = signal.signal(signal.SIGINT, handler) + except ValueError: + self._original_sigint_handler = None + + def _restore_signal_handler(self) -> None: + """Restore original SIGINT handler.""" + if self._original_sigint_handler is not None: + try: + signal.signal(signal.SIGINT, self._original_sigint_handler) + except ValueError: + pass + async def __run(self) -> CrawlResult: token = set_logger(self.logger) try: - engine = CrawlerEngine(self, session_manager=self._session_manager) - stats = await engine.crawl() - return CrawlResult(stats=stats, items=engine.items) + self._engine = CrawlerEngine(self, self._session_manager, self.crawldir, self._interval) + stats = await self._engine.crawl() + paused = self._engine.paused + return CrawlResult(stats=stats, items=self._engine.items, paused=paused) finally: + self._engine = None reset_logger(token) # Close any file handlers to release file resources. if self.log_file: @@ -213,27 +262,41 @@ class Spider(ABC): This is the main entry point for running a spider. Handles async execution internally via anyio. + If crawldir is set, pressing Ctrl+C will pause the spider and save a checkpoint. + Running the spider again with the same crawldir will resume from the checkpoint. + :param use_uvloop: Whether to use the faster uvloop/winloop event loop implementation, if available. :param backend_options: Asyncio backend options to be used with `anyio.run` """ backend_options = backend_options or {} if use_uvloop: backend_options.update({"use_uvloop": True}) - return anyio.run(self.__run, backend="asyncio", backend_options=backend_options) + + # Set up SIGINT handler for graceful pause (only if crawldir is set) + if self.crawldir: + self._setup_signal_handler() + + try: + return anyio.run(self.__run, backend="asyncio", backend_options=backend_options) + finally: + if self.crawldir: + self._restore_signal_handler() async def stream(self) -> AsyncGenerator[Dict[str, Any], None]: """Stream items as they're scraped. Ideal for long-running spiders or building applications on top of the spiders. Must be called from an async context. Yields items one by one as they are scraped. Access `spider.stats` during iteration for real-time statistics. + + Note: SIGINT handling for pause/resume is not available in stream mode. """ token = set_logger(self.logger) try: - self._stream_engine = CrawlerEngine(self, self._session_manager) - async for item in self._stream_engine: + self._engine = CrawlerEngine(self, self._session_manager, self.crawldir, self._interval) + async for item in self._engine: yield item finally: - self._stream_engine = None + self._engine = None reset_logger(token) if self.log_file: for handler in self.logger.handlers: @@ -243,6 +306,6 @@ class Spider(ABC): @property def stats(self) -> CrawlStats: """Access current crawl stats (works during streaming).""" - if self._stream_engine: - return self._stream_engine.stats + if self._engine: + return self._engine.stats raise RuntimeError("No active crawl. Use this property inside `async for item in spider.stream():`")