feat(spiders): Add pause/resume system for crawls

This commit is contained in:
Karim shoair
2026-01-19 03:43:10 +02:00
parent acdb39988c
commit 1f0f86e3a2
6 changed files with 336 additions and 28 deletions
+90
View File
@@ -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}")
+99 -8
View File
@@ -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()
+27 -1
View File
@@ -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
+6
View File
@@ -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)
+34 -2
View File
@@ -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")
+80 -17
View File
@@ -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():`")