diff --git a/compass/pipeline/collection/steps.py b/compass/pipeline/collection/steps.py index c060040b3..e2a7dc5ab 100644 --- a/compass/pipeline/collection/steps.py +++ b/compass/pipeline/collection/steps.py @@ -1,5 +1,6 @@ """Fixed collection steps for the process pipeline""" +import asyncio import logging from abc import ABC, abstractmethod @@ -267,22 +268,35 @@ async def collect(self, workflow): # ruff:ignore[no-self-use] ) return [] + crawl_timeout_s = ( + workflow.runtime.search_params.website_crawl_timeout_seconds + ) logger.debug( "Collecting documents using ELM web crawl for: %s", workflow.jurisdiction.full_name, ) try: - out = await download_jurisdiction_ordinances_from_website( - workflow.jurisdiction_website, - heuristic=await workflow.extractor.get_heuristic(), - keyword_points=( - await workflow.extractor.get_website_keywords() - ), - file_loader_kwargs=workflow.runtime.file_loader_kwargs, - crawl_semaphore=workflow.runtime.crawl_semaphore, - pb_jurisdiction_name=workflow.jurisdiction.full_name, - return_c4ai_results=True, + async with workflow.runtime.crawl_semaphore: + async with asyncio.timeout(crawl_timeout_s): + out = await download_jurisdiction_ordinances_from_website( + workflow.jurisdiction_website, + heuristic=await workflow.extractor.get_heuristic(), + keyword_points=( + await workflow.extractor.get_website_keywords() + ), + file_loader_kwargs=workflow.runtime.file_loader_kwargs, + pb_jurisdiction_name=workflow.jurisdiction.full_name, + return_c4ai_results=True, + ) + except TimeoutError: + logger.exception( + "ELM Website crawl deadline (%s) exceeded for %s; " + "continuing with no crawl docs", + f"{int(crawl_timeout_s):,d}s", + workflow.jurisdiction.full_name, ) + workflow.last_scrape_results = [] + return [] except Exception: logger.exception( "Error collecting documents using ELM web crawl for %s", @@ -358,17 +372,32 @@ async def collect(self, workflow): # ruff:ignore[no-self-use] for scrape_result in workflow.last_scrape_results: checked_urls.update({sub_res.url for sub_res in scrape_result}) + crawl_timeout_s = ( + workflow.runtime.search_params.website_crawl_timeout_seconds + ) func = download_jurisdiction_ordinances_from_website_compass_crawl try: - docs = await func( - workflow.jurisdiction_website, - heuristic=await workflow.extractor.get_heuristic(), - keyword_points=await workflow.extractor.get_website_keywords(), - file_loader_kwargs=workflow.runtime.file_loader_kwargs, - already_visited=checked_urls, - crawl_semaphore=workflow.runtime.crawl_semaphore, - pb_jurisdiction_name=workflow.jurisdiction.full_name, + async with workflow.runtime.crawl_semaphore: + async with asyncio.timeout(crawl_timeout_s): + docs = await func( + workflow.jurisdiction_website, + heuristic=await workflow.extractor.get_heuristic(), + keyword_points=( + await workflow.extractor.get_website_keywords() + ), + file_loader_kwargs=workflow.runtime.file_loader_kwargs, + already_visited=checked_urls, + pb_jurisdiction_name=workflow.jurisdiction.full_name, + ) + except TimeoutError: + logger.exception( + "COMPASS Website crawl deadline (%s) exceeded for %s; " + "continuing with no crawl docs", + f"{int(crawl_timeout_s):,d}s", + workflow.jurisdiction.full_name, ) + workflow.last_scrape_results = [] + return [] except Exception: logger.exception( "Error collecting documents using COMPASS web crawl for %s", diff --git a/compass/pipeline/data_classes.py b/compass/pipeline/data_classes.py index 2c521ff5a..9f3afbf47 100644 --- a/compass/pipeline/data_classes.py +++ b/compass/pipeline/data_classes.py @@ -181,6 +181,7 @@ def __init__( num_urls_to_check_per_jurisdiction=5, max_num_concurrent_browsers=10, max_num_concurrent_website_searches=None, + website_crawl_timeout_seconds=3600, url_ignore_substrings=None, url_keep_substrings=None, search_engines=None, @@ -205,6 +206,11 @@ def __init__( simultaneously. Increasing this value can speed up searches, but may lead to timeouts or performance issues on machines with limited resources. By default, ``10``. + website_crawl_timeout_seconds : int, default=3600 + Maximum number of seconds to allow for a website crawl to + complete before timing out. If the crawl exceeds this time, + it will be terminated and no documents will be returned for + the crawl step for the jurisdiction. By default, ``3600`` url_ignore_substrings : list of str, optional A list of substrings that, if found in any URL, will cause the URL to be excluded from consideration. This can be used @@ -273,6 +279,7 @@ def __init__( self.max_num_concurrent_website_searches = ( max_num_concurrent_website_searches ) + self.website_crawl_timeout_seconds = website_crawl_timeout_seconds self.url_ignore_substrings = _DOMAINS["blacklist"] self.url_ignore_substrings += url_ignore_substrings or [] self.url_keep_substrings = _DOMAINS["whitelist"] @@ -334,6 +341,7 @@ def __init__( # ruff:ignore[too-many-arguments] max_num_concurrent_browsers=10, max_num_concurrent_website_searches=10, max_num_concurrent_jurisdictions=25, + website_crawl_timeout_seconds=3600, url_ignore_substrings=None, url_keep_substrings=None, known_local_docs=None, @@ -474,6 +482,11 @@ def __init__( # ruff:ignore[too-many-arguments] Maximum number of jurisdictions to process concurrently. Limiting this can help manage memory usage when dealing with a large number of documents. By default, ``25``. + website_crawl_timeout_seconds : int, default=3600 + Maximum number of seconds to allow for a website crawl to + complete before timing out. If the crawl exceeds this time, + it will be terminated and no documents will be returned for + the crawl step for the jurisdiction. By default, ``3600`` url_ignore_substrings : list of str, optional A list of substrings that, if found in any URL, will cause the URL to be excluded from consideration. This can be used @@ -634,6 +647,7 @@ def __init__( # ruff:ignore[too-many-arguments] max_num_concurrent_website_searches=( max_num_concurrent_website_searches ), + website_crawl_timeout_seconds=website_crawl_timeout_seconds, url_ignore_substrings=url_ignore_substrings, url_keep_substrings=url_keep_substrings, search_engines=search_engines, @@ -702,6 +716,7 @@ def __init__( # ruff:ignore[too-many-arguments] max_num_concurrent_browsers=10, max_num_concurrent_website_searches=10, max_num_concurrent_jurisdictions=25, + website_crawl_timeout_seconds=3600, url_ignore_substrings=None, url_keep_substrings=None, known_local_docs=None, @@ -819,6 +834,11 @@ def __init__( # ruff:ignore[too-many-arguments] Maximum number of jurisdictions to process concurrently. Limiting this can help manage memory usage when dealing with a large number of documents. By default, ``25``. + website_crawl_timeout_seconds : int, default=3600 + Maximum number of seconds to allow for a website crawl to + complete before timing out. If the crawl exceeds this time, + it will be terminated and no documents will be returned for + the crawl step for the jurisdiction. By default, ``3600`` url_ignore_substrings : list of str, optional A list of substrings that, if found in any URL, will cause the URL to be excluded from consideration. This can be used @@ -996,6 +1016,7 @@ def __init__( # ruff:ignore[too-many-arguments] max_num_concurrent_website_searches ), max_num_concurrent_jurisdictions=max_num_concurrent_jurisdictions, + website_crawl_timeout_seconds=website_crawl_timeout_seconds, url_ignore_substrings=url_ignore_substrings, url_keep_substrings=url_keep_substrings, known_local_docs=known_local_docs, diff --git a/compass/pipeline/runtime.py b/compass/pipeline/runtime.py index ad7b75070..7a3996c42 100644 --- a/compass/pipeline/runtime.py +++ b/compass/pipeline/runtime.py @@ -116,14 +116,21 @@ def browser_semaphore(self): ) @cached_property - def crawl_semaphore(self): - """Crawl concurrency limiter""" + def _crawl_semaphore(self): + """Crawl concurrency limiter or None""" if not self.search_params.max_num_concurrent_website_searches: return None return asyncio.Semaphore( self.search_params.max_num_concurrent_website_searches ) + @property + def crawl_semaphore(self): + """Crawl concurrency limiter""" + if self._crawl_semaphore is None: + return AsyncExitStack() + return self._crawl_semaphore + @cached_property def search_engine_semaphore(self): """Search engine concurrency limiter""" @@ -172,6 +179,9 @@ def local_file_loader_kwargs(self): "html_read_kwargs": self.file_loader_kwargs.get( "html_read_kwargs" ), + "pdf_pipeline_options": self.file_loader_kwargs.get( + "pdf_pipeline_options" + ), } if self.search_params.pytesseract_exe_fp is not None: self._setup_pytesseract() @@ -240,7 +250,9 @@ def _base_services(self): ) if self.search_params.pytesseract_exe_fp is not None: - services.append(OCRPDFLoader(max_workers=1)) + kwargs = deepcopy(runtime_settings.ppe_kwargs or {}) + kwargs["max_workers"] = 1 + services.append(OCRPDFLoader(**kwargs)) return services @cached_property diff --git a/compass/plugin/one_shot/base.py b/compass/plugin/one_shot/base.py index 59b683372..4c7c8d011 100644 --- a/compass/plugin/one_shot/base.py +++ b/compass/plugin/one_shot/base.py @@ -590,14 +590,12 @@ def _out_cols_from_config(config): ) source_col_ind = next( - (ind for ind, col in enumerate(cols) if col.name == "source"), - None, + (ind for ind, col in enumerate(cols) if col.name == "source"), None ) - year_col = OutputColumn("year") if source_col_ind is None: - cols.append(year_col) + cols.extend((OutputColumn("year"), OutputColumn("source"))) else: - cols.insert(source_col_ind, year_col) + cols.insert(source_col_ind, OutputColumn("year")) cols.append( OutputColumn( diff --git a/compass/scripts/download.py b/compass/scripts/download.py index e821530bd..e915a5762 100644 --- a/compass/scripts/download.py +++ b/compass/scripts/download.py @@ -272,7 +272,6 @@ async def download_jurisdiction_ordinances_from_website( browser_config_kwargs=None, crawler_config_kwargs=None, max_urls=100, - crawl_semaphore=None, pb_jurisdiction_name=None, return_c4ai_results=False, ): @@ -307,10 +306,6 @@ async def download_jurisdiction_ordinances_from_website( max_urls : int, optional Max number of URLs to check from the website before terminating the search. By default, ``100``. - crawl_semaphore : :class:`asyncio.Semaphore`, optional - Semaphore instance that can be used to limit the number of - website searches happening concurrently. If ``None``, no limits - are applied. By default, ``None``. pb_jurisdiction_name : str, optional Optional jurisdiction name to use to update progress bar, if it's being used. By default, ``None``. @@ -337,9 +332,6 @@ async def download_jurisdiction_ordinances_from_website( to be running. """ - if crawl_semaphore is None: - crawl_semaphore = AsyncExitStack() - async def _doc_heuristic(doc): # ruff:ignore[unused-async] """Heuristic check for wind ordinance documents""" is_valid_document = heuristic.check(doc.text.lower()) @@ -393,7 +385,7 @@ async def _crawl_hook(*__, **___): # ruff:ignore[unused-async] cpb = AsyncExitStack() ch = None - async with crawl_semaphore, cpb: + async with cpb: docs_or_pair = await crawler.run( website, on_result_hook=ch, @@ -416,7 +408,6 @@ async def download_jurisdiction_ordinances_from_website_compass_crawl( already_visited=None, num_link_scores_to_check_per_page=4, max_urls=100, - crawl_semaphore=None, pb_jurisdiction_name=None, ): """Download ord documents from a website using the COMPASS crawler @@ -452,10 +443,6 @@ async def download_jurisdiction_ordinances_from_website_compass_crawl( max_urls : int, default=100 Max number of URLs to check from the website before terminating the search. By default, ``100``. - crawl_semaphore : :class:`asyncio.Semaphore`, optional - Semaphore instance that can be used to limit the number of - website crawls happening concurrently. If ``None``, no limits - are applied. By default, ``None``. pb_jurisdiction_name : str, optional Optional jurisdiction name to use to update progress bar, if it's being used. By default, ``None``. @@ -472,8 +459,6 @@ async def download_jurisdiction_ordinances_from_website_compass_crawl( Requires :class:`~compass.services.threaded.TempFileCache` service to be running. """ - if crawl_semaphore is None: - crawl_semaphore = AsyncExitStack() async def _doc_heuristic(doc): # ruff:ignore[unused-async] """Heuristic check for wind ordinance documents""" @@ -515,7 +500,7 @@ async def _crawl_hook(*__, **___): # ruff:ignore[unused-async] cpb = AsyncExitStack() ch = None - async with crawl_semaphore, cpb: + async with cpb: return await crawler.run(website, on_new_page_visit_hook=ch) diff --git a/compass/services/cpu.py b/compass/services/cpu.py index 057d1d3bc..2080f8935 100644 --- a/compass/services/cpu.py +++ b/compass/services/cpu.py @@ -2,8 +2,9 @@ import ast import os -import sys import time +import math +import pprint import asyncio import logging import warnings @@ -15,7 +16,6 @@ from pathlib import Path from functools import partial from concurrent.futures import ProcessPoolExecutor -from logging.handlers import QueueHandler import numpy as np import pandas as pd @@ -35,19 +35,26 @@ ) from docling.exceptions import ConversionError +from compass.exceptions import COMPASSValueError from compass.services.base import Service -from compass.utilities.logs import AddLocationFilter, LQ +from compass.utilities.logs import ( + configure_subprocess_logging, + configure_docling_subprocess_logging, + LQ, +) logger = logging.getLogger(__name__) +TIMEOUT_PARAMS = { + "shutdown_timeout": 5, + "force_shutdown_timeout": 1, +} class ProcessPoolService(Service): """Service that contains a ProcessPoolExecutor instance""" _DEFAULT_MAX_TASKS_PER_CHILD = 100 - _SHUTDOWN_TIMEOUT = 5 - _FORCE_SHUTDOWN_TIMEOUT = 1 def __init__(self, **kwargs): """ @@ -67,10 +74,11 @@ def acquire_resources(self): os.environ.setdefault("OMP_NUM_THREADS", "1") ppe_kwargs = dict(self._ppe_kwargs) # ppe_kwargs = self._set_tasks_per_child(ppe_kwargs) - user_initializer = ppe_kwargs.pop("initializer", None) - initargs = tuple(ppe_kwargs.pop("initargs", ())) - ppe_kwargs["initializer"] = _configure_subprocess_logging - ppe_kwargs["initargs"] = (LQ.QUEUE, user_initializer, initargs) + ppe_kwargs = self._set_ppe_initializer(ppe_kwargs) + logger.debug( + " - Setting up ProcessPoolExecutor with kwargs:\n%s", + pprint.PrettyPrinter().pformat(ppe_kwargs), + ) self.pool = ProcessPoolExecutor(**ppe_kwargs) def _set_tasks_per_child(self, ppe_kwargs): @@ -83,6 +91,14 @@ def _set_tasks_per_child(self, ppe_kwargs): ) return ppe_kwargs + def _set_ppe_initializer(self, ppe_kwargs): # ruff:ignore[no-self-use] + """Set initializer to configure subprocess logging""" + user_initializer = ppe_kwargs.pop("initializer", None) + initargs = tuple(ppe_kwargs.pop("initargs", ())) + ppe_kwargs["initializer"] = configure_subprocess_logging + ppe_kwargs["initargs"] = (LQ.QUEUE, user_initializer, initargs) + return ppe_kwargs + def release_resources(self): """Shutdown thread pool and cleanup temp directory""" pool = self.pool @@ -95,20 +111,20 @@ def release_resources(self): pool.shutdown(wait=False, cancel_futures=True) if not _needs_forced_shutdown( - manager_thread, processes, self._SHUTDOWN_TIMEOUT + manager_thread, processes, TIMEOUT_PARAMS["shutdown_timeout"] ): return logger.warning( "Process pool did not shut down within %.1f seconds; " "terminating lingering workers", - self._SHUTDOWN_TIMEOUT, + TIMEOUT_PARAMS["shutdown_timeout"], ) _force_shutdown_processes( - processes, timeout=self._FORCE_SHUTDOWN_TIMEOUT + processes, timeout=TIMEOUT_PARAMS["force_shutdown_timeout"] ) _join_manager_thread( - manager_thread, timeout=self._FORCE_SHUTDOWN_TIMEOUT + manager_thread, timeout=TIMEOUT_PARAMS["force_shutdown_timeout"] ) @@ -251,7 +267,9 @@ async def read_pdf_file_ocr(pdf_fp, **kwargs): ) -async def read_docling_web_file(doc_bytes, url, source_uri=None, **kwargs): +async def read_docling_web_file( + doc_bytes, url, source_uri=None, pdf_pipeline_options=None, **kwargs +): """Read a web file using Docling in a Process Pool Parameters @@ -264,6 +282,12 @@ async def read_docling_web_file(doc_bytes, url, source_uri=None, **kwargs): Original remote URL for the file. If specified, this is used as the HTML base URI while ``url`` is still used as the stream name for Docling format inference. By default, ``None``. + pdf_pipeline_options : dict, optional + Dictionary of keyword-value arguments to pass to + :class:`docling.datamodel.pipeline_options.PdfPipelineOptions` + initializer. Table structure defaults to enabled when omitted. + OCR defaults to enabled only when ``pytesseract_exe_fp`` is + specified. By default, ``None``. **kwargs Additional keyword arguments passed to Docling's :func:`~docling_core.types.doc.DoclingDocument.export_to_markdown` @@ -279,6 +303,7 @@ async def read_docling_web_file(doc_bytes, url, source_uri=None, **kwargs): doc_bytes, file_source=url, source_uri=source_uri, + pdf_pipeline_options=pdf_pipeline_options, **kwargs, ) @@ -357,9 +382,10 @@ def _read_docling_catch_error( headers=None, pytesseract_exe_fp=None, source_uri=None, + pdf_pipeline_options=None, **kwargs, ): - """Utility to return empty docs on Docling conversion errors""" + """Utility to return empty docs on error""" try: return _read_docling( doc_bytes=doc_bytes, @@ -367,9 +393,10 @@ def _read_docling_catch_error( headers=headers, pytesseract_exe_fp=pytesseract_exe_fp, source_uri=source_uri, + pdf_pipeline_options=pdf_pipeline_options, **kwargs, ) - except ConversionError: + except Exception: # ruff:ignore[blind-except] return MDDocument(pages=[], attrs={"doc_type": "unknown"}) @@ -379,16 +406,47 @@ def _read_docling( headers=None, pytesseract_exe_fp=None, source_uri=None, + pdf_pipeline_options=None, **kwargs, ): - """Utility func to read documents using Docling""" + """Read documents using Docling with an optional hard deadline""" + + pdf_pipeline_options = dict(pdf_pipeline_options or {}) + docling_timeout = pdf_pipeline_options.get( + "document_timeout", 60 * 60 / 1.1 + ) + _validate_docling_timeout(docling_timeout) + return _run_docling_in_subprocess( + _read_docling_without_timeout, + args=(doc_bytes, file_source), + kwargs={ + "headers": headers, + "pytesseract_exe_fp": pytesseract_exe_fp, + "source_uri": source_uri, + "pdf_pipeline_options": pdf_pipeline_options, + **kwargs, + }, + timeout=docling_timeout * 1.1, + ) + + +def _read_docling_without_timeout( + doc_bytes, + file_source, + headers=None, + pytesseract_exe_fp=None, + source_uri=None, + pdf_pipeline_options=None, + **kwargs, +): + """Read documents using Docling without an in-process deadline""" file_source = str(file_source) source_uri = file_source if source_uri is None else str(source_uri) if headers is not None: headers = dict(headers) - pipeline_options = PdfPipelineOptions() + pipeline_options = PdfPipelineOptions(**(pdf_pipeline_options or {})) pipeline_options.do_table_structure = True pipeline_options.table_structure_options = TableStructureOptions( do_cell_matching=True @@ -430,6 +488,7 @@ def _read_docling( "doc_filename": conv_result.input.file.stem, "doc_type": conv_result.input.format.value, "conversion_time_seconds": conversion_time_seconds, + "conversion_status": conv_result.status.value, "num_pages": len(conv_result.pages), "from_ocr": any( ~np.isnan(c.ocr_score) @@ -443,6 +502,96 @@ def _read_docling( return MDDocument([doc_text], attrs=attrs, remove_comments=False) +def _run_docling_in_subprocess(fn, *, args, kwargs, timeout): + """Run one Docling conversion in a disposable child process""" + mp_context = multiprocessing.get_context( + "fork" + if "fork" in multiprocessing.get_all_start_methods() + else "spawn" + ) + receiver, sender = mp_context.Pipe(duplex=False) + process = mp_context.Process( + target=_run_docling_subprocess, args=(sender, fn, args, kwargs) + ) + process.start() + sender.close() + + try: + status, payload = _receive_docling_result(receiver, process, timeout) + finally: + receiver.close() + _shutdown_docling_process(process) + + if status == "success": + logger.info("Docling conversion ran successfully in subprocess") + return payload + + msg = f"Docling conversion subprocess failed: {payload}" + logger.error(msg) + raise ConversionError(msg) + + +def _run_docling_subprocess(sender, fn, args, kwargs): + """Execute a Docling conversion and send its result to the worker""" + configure_docling_subprocess_logging(sender) + try: + sender.send(("success", fn(*args, **kwargs))) + except Exception as error: # ruff:ignore[blind-except] + with contextlib.suppress(BrokenPipeError, EOFError, OSError): + sender.send(("error", f"{type(error).__name__}: {error}")) + finally: + sender.close() + + +def _receive_docling_result(receiver, process, timeout): + """Receive a child conversion result before its deadline expires""" + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + msg = f"Docling conversion exceeded {timeout} seconds" + raise TimeoutError(msg) + + if receiver.poll(min(remaining, 10)): + try: + status, payload = receiver.recv() + except EOFError as error: + msg = "Docling conversion subprocess exited without a result" + raise ConversionError(msg) from error + + if status == "log": + logger.handle(payload) + continue + return status, payload + + if not process.is_alive(): + if receiver.poll(): + continue + msg = "Docling conversion subprocess exited without a result" + raise ConversionError(msg) + + +def _shutdown_docling_process(process): + """Join or force-stop a completed or timed-out Docling process""" + process.join(timeout=TIMEOUT_PARAMS["shutdown_timeout"]) + if process.is_alive(): + _force_shutdown_processes( + [process], timeout=TIMEOUT_PARAMS["force_shutdown_timeout"] + ) + + +def _validate_docling_timeout(timeout): + """Validate the configured Docling deadline""" + if ( + isinstance(timeout, bool) + or not isinstance(timeout, (int, float)) + or not math.isfinite(timeout) + or timeout <= 0 + ): + msg = "`docling_timeout` must be a positive number of seconds" + raise COMPASSValueError(msg) + + def _read_file_docling(fp, **kwargs): """Read a local file using Docling""" @@ -471,7 +620,8 @@ def _pytesseract_cleanup_win(temp_name): patches cleanup to suppress all OSErrors so the OCR result is not lost. """ - for filename in iglob(f"{temp_name}*" if temp_name else temp_name): # ruff:ignore[glob] + # ruff:ignore[glob] + for filename in iglob(f"{temp_name}*" if temp_name else temp_name): with contextlib.suppress(OSError): os.remove(filename) # ruff:ignore[os-remove] @@ -486,7 +636,8 @@ def _try_decode_ocr_pages(pages): decoded_pages = [] for page in pages: with contextlib.suppress(Exception): - page = ast.literal_eval(page).decode("utf-8") # ruff:ignore[redefined-loop-name] + # ruff:ignore[redefined-loop-name] + page = ast.literal_eval(page).decode("utf-8") decoded_pages.append(page) return decoded_pages @@ -535,71 +686,3 @@ def _force_shutdown_processes(processes, timeout=1): def _is_process_alive(process): """bool: Check whether a worker process is still alive""" return process is not None and process.is_alive() - - -def _configure_subprocess_logging(logging_queue, user_initializer, initargs): - """Route subprocess output through the main process log queue""" - queue_handler = QueueHandler(logging_queue) - queue_handler.addFilter(AddLocationFilter()) - - root_logger = logging.getLogger() - root_logger.handlers = [] - root_logger.addHandler(queue_handler) # root emits to queue handler - root_logger.setLevel(logging.INFO) - - for lib in ("compass", "elm", "docling", "openai"): - lib_logger = logging.getLogger(lib) - lib_logger.handlers = [] # no handlers within subprocess - lib_logger.propagate = True # instead, propogate to root logger - lib_logger.setLevel(logging.INFO) - - stdout_logger = logging.getLogger("compass.subprocess.stdout") - stderr_logger = logging.getLogger("compass.subprocess.stderr") - stdout_logger.setLevel(logging.INFO) - stderr_logger.setLevel(logging.WARNING) - sys.stdout = _LogStream(stdout_logger, logging.INFO) - sys.stderr = _LogStream(stderr_logger, logging.WARNING) - - if user_initializer is not None: - user_initializer(*initargs) - - -class _LogStream: - """File-like object that forwards writes into a logger""" - - def __init__(self, logger, level): - """ - - Parameters - ---------- - logger : logging.Logger - Logger to emit redirected stream output to. - level : int - Logging level used for forwarded messages. - """ - self.logger = logger - self.level = level - self._buffer = "" - self.encoding = "utf-8" - - def write(self, message): - """Forward complete lines to the configured logger""" - if not message: - return 0 - - self._buffer += message - while "\n" in self._buffer: - line, self._buffer = self._buffer.split("\n", 1) - if line: - self.logger.log(self.level, line) - return len(message) - - def flush(self): - """Flush any partial line buffered from the stream""" - if self._buffer: - self.logger.log(self.level, self._buffer) - self._buffer = "" - - def isatty(self): # ruff:ignore[no-self-use] - """bool: Redirected subprocess streams are never TTYs""" - return False diff --git a/compass/services/provider.py b/compass/services/provider.py index b2f16bbdc..71315d590 100644 --- a/compass/services/provider.py +++ b/compass/services/provider.py @@ -133,9 +133,9 @@ async def __aenter__(self): logger.debug("Initializing Service: %s", service.name) with contextlib.suppress(AttributeError): logger.debug( - " ↪ model_name=%r, rate_limit=%d", + " ↪ model_name=%r, rate_limit=%s", service.model_name, - service.rate_limit, + f"{int(service.rate_limit):,d}", ) queue = initialize_service_queue(service.name) service.acquire_resources() diff --git a/compass/services/threaded.py b/compass/services/threaded.py index 465dd205d..4fe6c2d58 100644 --- a/compass/services/threaded.py +++ b/compass/services/threaded.py @@ -17,6 +17,7 @@ from elm.web.utilities import write_url_doc_to_file from compass.services.base import Service +from compass.utilities.io import normalize_output_stem from compass.utilities.parsing import is_pdf_doc from compass.utilities import compute_cost_from_totals from compass.pb import COMPASS_PB @@ -51,19 +52,9 @@ def _compute_sha256(file_path): return f"sha256:{m.hexdigest()}" -def _normalize_output_stem(out_stem): - """Normalize an output file name while preserving the full stem""" - return ( - out_stem.replace(".", "") - .replace(",", "") - .replace("/", "_") - .replace(" ", "_") - ) - - def _ensure_output_suffix(out_dir, out_stem, suffix): """Build output path""" - out_stem = _normalize_output_stem(out_stem) + out_stem = normalize_output_stem(out_stem) return Path(out_dir) / f"{out_stem}{suffix}" diff --git a/compass/utilities/costs.py b/compass/utilities/costs.py index fa81f9d7c..47f9bba29 100644 --- a/compass/utilities/costs.py +++ b/compass/utilities/costs.py @@ -17,6 +17,9 @@ "gpt-5.4-mini": {"prompt": 0.75, "response": 4.5}, "gpt-5.4-nano": {"prompt": 0.20, "response": 1.25}, "gpt-5.5": {"prompt": 5, "response": 30}, + "gpt-5.6-sol": {"prompt": 5, "response": 30}, + "gpt-5.6-terra": {"prompt": 2, "response": 12}, + "gpt-5.6-luna": {"prompt": 0.2, "response": 1.2}, "compassop-gpt-4o": {"prompt": 2.5, "response": 10}, "compassop-gpt-4o-mini": {"prompt": 0.15, "response": 0.6}, "compassop-gpt-4.1": {"prompt": 2, "response": 8}, diff --git a/compass/utilities/io.py b/compass/utilities/io.py index f2e166c8f..4a1682eb0 100644 --- a/compass/utilities/io.py +++ b/compass/utilities/io.py @@ -339,3 +339,14 @@ def resolve_path(path, base_dir): path = path.expanduser().resolve().as_posix() return path + + +def normalize_output_stem(out_stem): + """[NOT PUBLIC API] Normalize an output file name stem""" + return ( + out_stem.replace(".", "") + .replace(",", "") + .replace("/", "_") + .replace("\\", "_") + .replace(" ", "_") + ) diff --git a/compass/utilities/logs.py b/compass/utilities/logs.py index cd0ebef89..8e334352b 100644 --- a/compass/utilities/logs.py +++ b/compass/utilities/logs.py @@ -5,12 +5,14 @@ """ import os +import sys import time import json import copy import asyncio import logging import threading +import contextlib import multiprocessing from pathlib import Path from functools import partial, partialmethod @@ -18,6 +20,7 @@ from importlib.metadata import version, PackageNotFoundError from compass import __version__ +from compass.utilities.io import normalize_output_stem from compass.exceptions import COMPASSValueError @@ -28,8 +31,7 @@ class _LQ: """Logging queue descriptor""" def __get__(self, __, lq_class=None): - lq_class.QUEUE = multiprocessing.get_context().Queue() - # lq_class.QUEUE = multiprocessing.get_context("spawn").Queue() + lq_class.QUEUE = multiprocessing.get_context("spawn").Queue() return lq_class.QUEUE @@ -169,6 +171,47 @@ def emit(self, record): self.handleError(record) +class _LogStream: + """File-like object that forwards writes into a logger""" + + def __init__(self, logger, level): + """ + + Parameters + ---------- + logger : logging.Logger + Logger to emit redirected stream output to. + level : int + Logging level used for forwarded messages. + """ + self.logger = logger + self.level = level + self._buffer = "" + self.encoding = "utf-8" + + def write(self, message): + """Forward complete lines to the configured logger""" + if not message: + return 0 + + self._buffer += message + while "\n" in self._buffer: + line, self._buffer = self._buffer.split("\n", 1) + if line: + self.logger.log(self.level, line) + return len(message) + + def flush(self): + """Flush any partial line buffered from the stream""" + if self._buffer: + self.logger.log(self.level, self._buffer) + self._buffer = "" + + def isatty(self): # ruff:ignore[no-self-use] + """bool: Redirected subprocess streams are never TTYs""" + return False + + class LogListener: """Class to listen to logging queue and write logs to files""" @@ -316,8 +359,9 @@ def _create_log_dir(self): def _setup_handler(self): """Setup the file handler for this location""" + fn_stem = normalize_output_stem(self.location) self._handler = logging.FileHandler( - self.log_dir / f"{self.location}.log", encoding="utf-8" + self.log_dir / f"{fn_stem}.log", encoding="utf-8" ) self._handler.setLevel(self.level) self._handler.addFilter(LocationFilter(self.location)) @@ -325,8 +369,9 @@ def _setup_handler(self): def _setup_exception_handler(self): """Setup file handler for tracking errors for this location""" + fn_stem = normalize_output_stem(f"{self.location}_exceptions") self._exception_handler = _JsonExceptionFileHandler( - self.log_dir / f"{self.location} exceptions.json", encoding="utf-8" + self.log_dir / f"{fn_stem}.json", encoding="utf-8" ) self._exception_handler.addFilter(LocationFilter(self.location)) @@ -406,7 +451,8 @@ async def __aexit__(self, exc_type, exc, tb): class ExceptionOnlyFilter(logging.Filter): """Filter to only pass through Exception logging (errors)""" - def filter(self, record): # ruff:ignore[undocumented-public-method, no-self-use] + # ruff:ignore[undocumented-public-method, no-self-use] + def filter(self, record): return bool(record.exc_info or getattr(record, "exc_type", None)) @@ -417,7 +463,8 @@ def format(self, record): exc_info, exc_text = _extract_exc_info_from_record(record) message = record.getMessage() - if message and len(message) > 103: # ruff:ignore[magic-value-comparison] + # ruff:ignore[magic-value-comparison] + if message and len(message) > 103: message = message[:103] return { @@ -502,6 +549,18 @@ def _get_existing_records(self): return records +class _DoclingLogPipe: + """Queue-like object that sends log records through a connection""" + + def __init__(self, sender): + self.sender = sender + + def put_nowait(self, record): + """Send one prepared log record to the parent process""" + with contextlib.suppress(BrokenPipeError, EOFError, OSError): + self.sender.send(("log", record)) + + def log_versions(logger): """Log COMPASS and dependency package versions @@ -547,6 +606,40 @@ def setup_logging_levels(): logging.debug_to_file = partial(logging.log, logging.DEBUG_TO_FILE) +def configure_subprocess_logging(logging_queue, user_initializer, initargs): + """[NOT PUBLIC API] Route subprocess output through main queue""" + queue_handler = QueueHandler(logging_queue) + queue_handler.addFilter(AddLocationFilter()) + queue_handler.setFormatter(logging.Formatter("[%(name)s] %(message)s")) + + root_logger = logging.getLogger() + root_logger.handlers = [] + root_logger.addHandler(queue_handler) # root emits to queue handler + root_logger.setLevel(logging.INFO) + + for lib in ("compass", "elm", "docling", "openai"): + lib_logger = logging.getLogger(lib) + lib_logger.handlers = [] # no handlers within subprocess + lib_logger.propagate = True # instead, propagate to root logger + lib_logger.setLevel(logging.INFO) + + stdout_logger = logging.getLogger("compass.subprocess.stdout") + stderr_logger = logging.getLogger("compass.subprocess.stderr") + stdout_logger.setLevel(logging.INFO) + stderr_logger.setLevel(logging.WARNING) + sys.stdout = _LogStream(stdout_logger, logging.INFO) + sys.stderr = _LogStream(stderr_logger, logging.WARNING) + + logging.getLogger("compass").info("Subprocess logging initialized") + if user_initializer is not None: + user_initializer(*initargs) + + +def configure_docling_subprocess_logging(sender): + """[NOT PUBLIC API] Route docling subprocess output through main""" + configure_subprocess_logging(_DoclingLogPipe(sender), None, ()) + + def _get_version(pkg_name): """Get the version string for a package""" try: diff --git a/compass/web/file_loader.py b/compass/web/file_loader.py index 8ae6b3afb..3a4f01747 100644 --- a/compass/web/file_loader.py +++ b/compass/web/file_loader.py @@ -86,7 +86,8 @@ async def _fetch_doc(self, url): class AsyncDoclingWebFileLoader(BaseAsyncFileLoader): """Async web file loader using Docling""" - def __init__( # ruff:ignore[too-many-arguments, too-many-positional-arguments] + # ruff:ignore[too-many-arguments, too-many-positional-arguments] + def __init__( self, header_template=None, verify_ssl=True, @@ -100,7 +101,9 @@ def __init__( # ruff:ignore[too-many-arguments, too-many-positional-arguments] num_pw_html_retries=3, to_md_kwargs=None, pytesseract_exe_fp=None, - **__, # consume any extra kwargs + pdf_pipeline_options=None, + re_fetch_failed_with_elm=False, + **extra, ): """ @@ -163,6 +166,29 @@ def __init__( # ruff:ignore[too-many-arguments, too-many-positional-arguments] Path to the `pytesseract` executable. If specified, OCR will be used to extract text from scanned PDFs using Google's Tesseract. By default ``None``. + pdf_pipeline_options : dict, optional + Dictionary of keyword-value arguments to pass to + :class:`docling.datamodel.pipeline_options.PdfPipelineOptions` + initializer. Table structure defaults to enabled when not + explicitly specified. OCR defaults to enabled only when a + Tesseract executable path is provided. If ``None``, the + default options are used. + re_fetch_failed_with_elm : bool, default=False + Option to re-fetch failed sources using ELM's default + fetcher. This can be useful if Docling fails to parse a + document, but ELM's fetcher can still retrieve it. To make + sure this functions properly, be sure to specify + ``pdf_read_kwargs`` and ``pdf_read_coroutine``, in the + ``extra`` kwargs as you would for the elm-based + :class:`~elm.web.file_loader.AsyncWebFileLoader`. + + .. NOTE:: + + This is meant to be a _fast_ fallback option for the + longer Docling parse, so OCR PDF parsing is completely + disabled for the ELM fallback. + + By default, ``False``. """ super().__init__(file_cache_coroutine=file_cache_coroutine) self.content_fetcher = AsyncFetchWithRetry( @@ -180,6 +206,25 @@ def __init__( # ruff:ignore[too-many-arguments, too-many-positional-arguments] ) self.to_md_kwargs = to_md_kwargs or {} self.pytesseract_exe_fp = pytesseract_exe_fp + self.pdf_pipeline_options = pdf_pipeline_options + + self.failed_fetcher = None + if re_fetch_failed_with_elm: + self.failed_fetcher = AsyncWebFileLoader( + header_template=header_template, + verify_ssl=verify_ssl, + aget_kwargs=aget_kwargs, + pw_launch_kwargs=pw_launch_kwargs, + pdf_read_kwargs=extra.get("pdf_read_kwargs"), + html_read_kwargs=html_read_kwargs, + pdf_read_coroutine=extra.get("pdf_read_coroutine"), + html_read_coroutine=html_read_coroutine, + pdf_ocr_read_coroutine=None, + file_cache_coroutine=file_cache_coroutine, + browser_semaphore=browser_semaphore, + use_scrapling_stealth=use_scrapling_stealth, + num_pw_html_retries=num_pw_html_retries, + ) async def fetch_all(self, *sources): """Fetch documents for all requested sources. @@ -195,6 +240,12 @@ async def fetch_all(self, *sources): list List of parsed documents. """ + docs = await self._fetch_docs_with_docling(sources) + docs = await self._fetch_html_docs_again_using_playwright(docs) + return await self._maybe_fetch_failed_docs_with_elm(docs, sources) + + async def _fetch_docs_with_docling(self, sources): + """Fetch docs using Docling""" outer_task_name = asyncio.current_task().get_name() fetches = [ asyncio.create_task(self.fetch(source), name=outer_task_name) @@ -212,21 +263,73 @@ async def fetch_all(self, *sources): ] ), ) + return docs + async def _fetch_html_docs_again_using_playwright(self, docs): + """Fetch HTML docs using Playwright""" to_re_fetch = [ doc.attrs["source"] for doc in docs if doc.attrs["doc_type"].casefold() == "html" ] - if to_re_fetch: - logger.debug( - "Loading HTML with Playwright for %d source(s):\n%r", - len(to_re_fetch), - to_re_fetch, - ) - docs += await self.html_loader.fetch_all(*to_re_fetch) + if not to_re_fetch: + return docs + + logger.debug( + "Loading HTML with Playwright for %d source(s):\n%r", + len(to_re_fetch), + to_re_fetch, + ) + docs += await self.html_loader.fetch_all(*to_re_fetch) return docs + async def _maybe_fetch_failed_docs_with_elm(self, docs, sources): + """Fetch failed docs using ELM (if enabled)""" + if self.failed_fetcher is None: + return docs + + out_docs = [] + partial_fail_docs = {} + failed_searches = [] + for source in sources: + source_docs = [ + doc for doc in docs if doc.attrs["source"] == source + ] + if not source_docs: + failed_searches.append(source) + continue + + if len(source_docs) > 1: + out_docs.extend(source_docs) + continue + + doc = source_docs[0] + if doc.attrs.get("conversion_status") != "success": + failed_searches.append(source) + partial_fail_docs[source] = doc + else: + out_docs.append(doc) + + if not failed_searches: + return out_docs + + logger.debug( + "Re-fetching %d failed source(s) with ELM:\n%r", + len(failed_searches), + failed_searches, + ) + elm_docs = await self.failed_fetcher.fetch_all(*failed_searches) + + for elm_doc in elm_docs: + docling_doc = partial_fail_docs.get(elm_doc.attrs["source"]) + elm_doc_failed = elm_doc.empty or "cache_fn" not in elm_doc.attrs + if elm_doc_failed and docling_doc is not None: + out_docs.append(docling_doc) + else: + out_docs.append(elm_doc) + + return out_docs + async def _fetch_doc(self, url): """Fetch a doc using Docling""" @@ -239,18 +342,34 @@ async def _fetch_doc(self, url): resolved_filename = resolve_remote_filename( http_url=AnyHttpUrl(url), response_headers=dict(headers) ) - doc = await read_docling_web_file( - raw_content, - url=resolved_filename, - source_uri=url, - headers=dict(headers), - pytesseract_exe_fp=self.pytesseract_exe_fp, - **self.to_md_kwargs, - ) + logger.debug("Docling is starting content read from %r", url) + try: + doc = await read_docling_web_file( + raw_content, + url=resolved_filename, + source_uri=url, + headers=dict(headers), + pytesseract_exe_fp=self.pytesseract_exe_fp, + pdf_pipeline_options=self.pdf_pipeline_options, + **self.to_md_kwargs, + ) + except TimeoutError: + logger.info("Docling parsing timed out for %r", url) + return MDDocument(pages=[]), None + if doc.empty: logger.info("Docling could not parse content from %s", url) return doc, None + logger.debug( + "Docling finished parsing %r:\n\t- Status: %r\n\t- " + "Conversion time (s): %.2f\n\t- Num pages: %r\n\t- From OCR: %r", + url, + doc.attrs.get("conversion_status", "unknown"), + doc.attrs.get("conversion_time_seconds", -1), + doc.attrs.get("num_pages", "unknown"), + doc.attrs.get("from_ocr", "unknown"), + ) if doc.attrs["doc_type"].casefold() != "html": doc.WRITE_KWARGS = {"mode": "wb"} doc.FILE_EXTENSION = doc.attrs["doc_type"] @@ -268,6 +387,7 @@ def __init__( doc_attrs=None, to_md_kwargs=None, pytesseract_exe_fp=None, + pdf_pipeline_options=None, **__, # consume any extra kwargs ): """ @@ -296,23 +416,47 @@ def __init__( Path to the `pytesseract` executable. If specified, OCR will be used to extract text from scanned PDFs using Google's Tesseract. By default ``None``. + pdf_pipeline_options : dict, optional + Dictionary of keyword-value arguments to pass to + :class:`docling.datamodel.pipeline_options.PdfPipelineOptions` + initializer. Table structure defaults to enabled when not + explicitly specified. OCR defaults to enabled only when a + Tesseract executable path is provided. If ``None``, the + default options are used. """ super().__init__(file_cache_coroutine=file_cache_coroutine) self.to_md_kwargs = to_md_kwargs or {} self.doc_attrs = doc_attrs or {} self.pytesseract_exe_fp = pytesseract_exe_fp + self.pdf_pipeline_options = pdf_pipeline_options async def _fetch_doc(self, source): """Load a doc by reading file based on extension""" - doc, raw_content = await read_docling_local_file( - source, - pytesseract_exe_fp=self.pytesseract_exe_fp, - **self.to_md_kwargs, - ) + logger.debug("Docling is starting content read from %s", source) + try: + doc, raw_content = await read_docling_local_file( + source, + pytesseract_exe_fp=self.pytesseract_exe_fp, + pdf_pipeline_options=self.pdf_pipeline_options, + **self.to_md_kwargs, + ) + except TimeoutError: + logger.info("Docling parsing timed out for %s", source) + return MDDocument(pages=[]), None + if doc.empty: logger.info("Docling could not parse content from %s", source) return doc, None + logger.debug( + "Docling finished parsing %s:\n\t- Status: %r\n\t- " + "Conversion time (s): %.2f\n\t- Num pages: %r\n\t- From OCR: %r", + source, + doc.attrs.get("conversion_status", "unknown"), + doc.attrs.get("conversion_time_seconds", -1), + doc.attrs.get("num_pages", "unknown"), + doc.attrs.get("from_ocr", "unknown"), + ) if doc.attrs["doc_type"].casefold() != "html": doc.WRITE_KWARGS = {"mode": "wb"} doc.FILE_EXTENSION = doc.attrs["doc_type"] diff --git a/compass/web/website_crawl.py b/compass/web/website_crawl.py index ec703bcfb..2871d49c1 100644 --- a/compass/web/website_crawl.py +++ b/compass/web/website_crawl.py @@ -5,6 +5,7 @@ interface). """ +import math import logging import operator from collections import Counter @@ -464,16 +465,24 @@ async def _get_text_no_err(self, url): async def _get_text(self, url): """Get all html text from a page""" all_text = [] + timeout_ms = 180_000 # milliseconds pw_page_kwargs = { "intercept_routes": True, "ignore_https_errors": True, - "timeout": 60_0000, + "timeout": timeout_ms, } async with async_playwright() as p, self.browser_semaphore: browser = await p.chromium.launch(**self.pw_launch_kwargs) async with pw_page(browser, **pw_page_kwargs) as page: await page.goto(url) - await page.wait_for_load_state("networkidle", timeout=60_000) + logger.debug( + "Waiting up to %d min for '%s' to load...", + math.ceil(timeout_ms / 60_000), + url, + ) + await page.wait_for_load_state( + "networkidle", timeout=timeout_ms + ) all_text.append(await page.content()) all_text += await _get_text_from_all_locators(page) diff --git a/examples/execution_basics/README.rst b/examples/execution_basics/README.rst index 87a4d65a5..49adb8a6a 100644 --- a/examples/execution_basics/README.rst +++ b/examples/execution_basics/README.rst @@ -134,6 +134,26 @@ You can locate the executable path by running: Omit the ``pytesseract_exe_fp`` key to disable OCR functionality. +**Docling conversion deadline** +When using the Docling file loader, set ``docling_timeout`` in +``file_loader_kwargs.pdf_pipeline_options`` to apply a wall-clock deadline to each document. +COMPASS runs the conversion in a disposable child process and terminates that +child when the deadline expires. + +.. code-block:: json + + "pytesseract_exe_fp": "/path/to/tesseract", + "file_loader_kwargs": { + "pdf_pipeline_options": { + "do_ocr": true, + "do_table_structure": true + "docling_timeout": 1800, + } + } + +This setting applies only when ``COMPASS_FILE_LOAD_BACKEND=docling``. The deadline wrapper +does not retry a timed-out document. + Kitchen Sink Config ------------------- diff --git a/tests/python/integration/test_integrated.py b/tests/python/integration/test_integrated.py index bc0adcd1e..075264aa9 100644 --- a/tests/python/integration/test_integrated.py +++ b/tests/python/integration/test_integrated.py @@ -230,7 +230,9 @@ async def search_location_with_logs( for fp in log_files: text = fp.read_text() assert "A generic test log" in text - assert f"This location is {fp.stem!r}" in text + assert any( + f"This location is {loc!r}" in text for loc in test_locations + ) @pytest.mark.asyncio diff --git a/tests/python/unit/pipeline/test_pipeline_collection_steps.py b/tests/python/unit/pipeline/test_pipeline_collection_steps.py index 5a40d27a2..ec7962f5a 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection_steps.py +++ b/tests/python/unit/pipeline/test_pipeline_collection_steps.py @@ -2,6 +2,7 @@ from pathlib import Path from types import SimpleNamespace +from contextlib import AsyncExitStack import pytest @@ -52,12 +53,13 @@ def _build_workflow(*, website="https://example.com", models=None): "loader_mode": "ocr", }, file_loader_kwargs_no_ocr={"loader_mode": "no-ocr"}, - crawl_semaphore=None, + crawl_semaphore=AsyncExitStack(), browser_semaphore=None, search_engine_semaphore=None, search_params=SimpleNamespace( url_ignore_substrings=(), se_kwargs={}, + website_crawl_timeout_seconds=3600, ), models=models, ) diff --git a/tests/python/unit/pipeline/test_pipeline_orchestration.py b/tests/python/unit/pipeline/test_pipeline_orchestration.py index d5392b572..b204afa30 100644 --- a/tests/python/unit/pipeline/test_pipeline_orchestration.py +++ b/tests/python/unit/pipeline/test_pipeline_orchestration.py @@ -116,6 +116,28 @@ def test_known_local_docs_logs_missing_file(tmp_path, testing_log_file): ) +def test_runtime_passes_docling_pipeline_options_to_local_loader(tmp_path): + """Pass Docling configuration to known-local document loaders""" + request = CollectionRequest( + out_dir=tmp_path / "outputs", + tech="solar", + jurisdiction_fp=tmp_path / "jurisdictions.csv", + file_loader_kwargs={ + "pdf_pipeline_options": { + "document_timeout": 120, + "do_table_structure": True, + }, + }, + ) + + runtime = PipelineRuntime(request) + + assert runtime.local_file_loader_kwargs["pdf_pipeline_options"] == { + "document_timeout": 120, + "do_table_structure": True, + } + + @pytest.mark.asyncio async def test_collect_request_uses_collection_workflow( tmp_path, patched_workflow diff --git a/tests/python/unit/services/test_services_cpu.py b/tests/python/unit/services/test_services_cpu.py index 231e1ae44..e7e204fe6 100644 --- a/tests/python/unit/services/test_services_cpu.py +++ b/tests/python/unit/services/test_services_cpu.py @@ -2,6 +2,7 @@ import logging import sys +import time import asyncio from pathlib import Path from types import SimpleNamespace @@ -10,7 +11,13 @@ import pandas as pd import pytest -from compass.services.cpu import ProcessPoolService, _read_docling +from compass.services.cpu import ( + TIMEOUT_PARAMS, + ProcessPoolService, + _read_docling, + _read_docling_without_timeout, + _run_docling_in_subprocess, +) from compass.services.provider import RunningAsyncServices from compass.utilities.logs import LocationFileLog, LogListener @@ -34,6 +41,16 @@ def _write_to_process_streams(): return "STREAMED" +def _return_from_subprocess(value): + """Return a serializable value from a child process""" + return value + + +def _block_subprocess(seconds): + """Block a child process long enough for a deadline to expire""" + time.sleep(seconds) + + @pytest.mark.asyncio async def test_logging_within_service(tmp_path): """Test that child-process logs are forwarded to the listener""" @@ -63,16 +80,21 @@ def emit(self, record): with LocationFileLog(ll, tmp_path, location="test_loc", level="DEBUG"): msg = await ProcessLogging.call() for _ in range(30): - if captured_records: + messages = {record.message for record in captured_records} + if "[compass] HELLO WORLD" in messages: break await asyncio.sleep(0.1) ll.removeHandler(capture_handler) assert msg == "HELLO WORLD" - assert any(record.message == "HELLO WORLD" for record in captured_records) + assert any( + record.message == "[compass] HELLO WORLD" + for record in captured_records + ), {record.message for record in captured_records} assert not any( - record.message == "A DEBUG LOG" for record in captured_records - ) + record.message == "[compass] A DEBUG LOG" + for record in captured_records + ), {record.message for record in captured_records} @pytest.mark.asyncio @@ -105,18 +127,24 @@ def emit(self, record): ll.addHandler(capture_handler) msg = await ProcessStreamLogging.call() for _ in range(30): - if len(captured_records) >= 2: + messages = {record.message for record in captured_records} + if { + "[compass.subprocess.stdout] PROCESS STDOUT", + "[compass.subprocess.stderr] PROCESS STDERR", + } <= messages: break await asyncio.sleep(0.1) ll.removeHandler(capture_handler) assert msg == "STREAMED" assert any( - record.message == "PROCESS STDOUT" for record in captured_records - ) + record.message == "[compass.subprocess.stdout] PROCESS STDOUT" + for record in captured_records + ), {record.message for record in captured_records} assert any( - record.message == "PROCESS STDERR" for record in captured_records - ) + record.message == "[compass.subprocess.stderr] PROCESS STDERR" + for record in captured_records + ), {record.message for record in captured_records} def test_read_docling_converts_missing_confidences_to_none(monkeypatch): @@ -139,21 +167,86 @@ def convert(self, stream, headers=None): ), pages=["page 1"], document=SimpleNamespace( + # ruff:ignore[unused-lambda-argument] export_to_markdown=lambda **kwargs: "markdown body" ), + status=SimpleNamespace(value="success"), ) monkeypatch.setattr( "compass.services.cpu.DocumentConverter", FakeDocumentConverter ) - doc = _read_docling(b"", "sample.html") + doc = _read_docling_without_timeout(b"", "sample.html") assert doc.pages == ["markdown body"] assert doc.attrs["mean_confidence"] is None assert doc.attrs["low_score_confidence"] is None +def test_read_docling_uses_process_deadline(monkeypatch): + """Docling deadlines should run outside the process-pool worker""" + captured = {} + expected = object() + configured_options = {"document_timeout": 120} + + def _run_in_subprocess(fn, *, args, kwargs, timeout): + captured["fn"] = fn + captured["args"] = args + captured["kwargs"] = kwargs + captured["timeout"] = timeout + return expected + + monkeypatch.setattr( + "compass.services.cpu._run_docling_in_subprocess", + _run_in_subprocess, + ) + + result = _read_docling( + b"%PDF", + "sample.pdf", + pdf_pipeline_options=configured_options, + ) + + assert result is expected + assert captured["fn"] is _read_docling_without_timeout + assert captured["args"] == (b"%PDF", "sample.pdf") + assert captured["kwargs"]["pdf_pipeline_options"] == { + "document_timeout": 120 + } + assert captured["timeout"] == 132 + assert configured_options == {"document_timeout": 120} + + +def test_docling_subprocess_returns_result(): + """The Docling child process should return completed conversions""" + result = _run_docling_in_subprocess( + _return_from_subprocess, + args=("converted",), + kwargs={}, + timeout=60, + ) + + assert result == "converted" + + +def test_docling_subprocess_enforces_deadline(monkeypatch): + """The Docling child process should be stopped at its deadline""" + monkeypatch.setitem(TIMEOUT_PARAMS, "shutdown_timeout", 0.1) + monkeypatch.setitem(TIMEOUT_PARAMS, "force_shutdown_timeout", 0.1) + + start_time = time.monotonic() + with pytest.raises(TimeoutError, match="Docling conversion exceeded"): + _run_docling_in_subprocess( + _block_subprocess, + args=(10,), + kwargs={}, + timeout=0.1, + ) + + assert time.monotonic() - start_time < 1 + + def test_process_pool_release_resources_graceful_shutdown(): # ruff:ignore[complex-structure] """Graceful process-pool shutdown should not force worker exit""" @@ -213,7 +306,7 @@ def shutdown(self, wait=True, cancel_futures=True): assert service.pool is None assert pool.shutdown_calls == [(False, True)] assert pool._executor_manager_thread.join_calls == [ - service._SHUTDOWN_TIMEOUT + TIMEOUT_PARAMS["shutdown_timeout"] ] process = pool._processes[0] assert process.terminate_calls == 0 @@ -282,15 +375,15 @@ def shutdown(self, wait=True, cancel_futures=True): assert service.pool is None assert pool.shutdown_calls == [(False, True)] assert pool._executor_manager_thread.join_calls == [ - service._SHUTDOWN_TIMEOUT, - service._FORCE_SHUTDOWN_TIMEOUT, + TIMEOUT_PARAMS["shutdown_timeout"], + TIMEOUT_PARAMS["force_shutdown_timeout"], ] process = pool._processes[0] assert process.terminate_calls == 1 assert process.kill_calls == 1 assert process.join_calls == [ - service._FORCE_SHUTDOWN_TIMEOUT, - service._FORCE_SHUTDOWN_TIMEOUT, + TIMEOUT_PARAMS["force_shutdown_timeout"], + TIMEOUT_PARAMS["force_shutdown_timeout"], ] diff --git a/tests/python/unit/utilities/test_utilities_io.py b/tests/python/unit/utilities/test_utilities_io.py index 26f81e2d3..f29c37fe5 100644 --- a/tests/python/unit/utilities/test_utilities_io.py +++ b/tests/python/unit/utilities/test_utilities_io.py @@ -81,26 +81,26 @@ def shutdown(self, wait=True, cancel_futures=True): # service.release_resources() -# def test_file_loader_preserves_max_tasks_per_child_override(monkeypatch): -# """Test process pool respects user task-recycling overrides""" +def test_file_loader_preserves_max_tasks_per_child_override(monkeypatch): + """Test process pool respects user task-recycling overrides""" -# captured_kwargs = {} + captured_kwargs = {} -# class DummyPool: -# def __init__(self, *__, **kwargs): -# captured_kwargs.update(kwargs) + class DummyPool: + def __init__(self, *__, **kwargs): + captured_kwargs.update(kwargs) -# def shutdown(self, wait=True, cancel_futures=True): -# return None + def shutdown(self, wait=True, cancel_futures=True): + return None -# monkeypatch.setattr("compass.services.cpu.ProcessPoolExecutor", DummyPool) + monkeypatch.setattr("compass.services.cpu.ProcessPoolExecutor", DummyPool) -# service = FileLoader(max_tasks_per_child=7) -# service.acquire_resources() + service = FileLoader(max_tasks_per_child=7) + service.acquire_resources() -# assert captured_kwargs["max_tasks_per_child"] == 7 + assert captured_kwargs["max_tasks_per_child"] == 7 -# service.release_resources() + service.release_resources() # def test_file_loader_sets_spawn_mp_context(monkeypatch): diff --git a/tests/python/unit/utilities/test_utilities_logs.py b/tests/python/unit/utilities/test_utilities_logs.py index 23faef3c5..0581bc48c 100644 --- a/tests/python/unit/utilities/test_utilities_logs.py +++ b/tests/python/unit/utilities/test_utilities_logs.py @@ -174,7 +174,7 @@ async def _produce_logs(listener): await task text_log = log_dir / "async_loc.log" - json_log = log_dir / "async_loc exceptions.json" + json_log = log_dir / "async_loc_exceptions.json" assert text_log.exists() assert json_log.exists() diff --git a/tests/python/unit/web/test_web_file_loader.py b/tests/python/unit/web/test_web_file_loader.py new file mode 100644 index 000000000..9893d9182 --- /dev/null +++ b/tests/python/unit/web/test_web_file_loader.py @@ -0,0 +1,98 @@ +"""COMPASS web file loader tests""" + +import asyncio +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from compass.web.file_loader import AsyncDoclingWebFileLoader + + +def _doc(source, doc_type="pdf", empty=False, conversion_status="success"): + return SimpleNamespace( + attrs={ + "source": source, + "doc_type": doc_type, + "conversion_status": conversion_status, + }, + empty=empty, + ) + + +class _Fetcher: + async def fetch(self, url): + return b"content", "application/pdf", None, {} + + +class _FailedFetcher: + def __init__(self, docs): + self.docs = docs + self.calls = [] + + async def fetch_all(self, *sources): + self.calls.append(sources) + return [self.docs[source] for source in sources] + + +@pytest.mark.asyncio +async def test_docling_web_file_loader_fetch_all_falls_back_to_elm( + monkeypatch, +): + """Retry only missing sources with the ELM fallback loader""" + loader = AsyncDoclingWebFileLoader() + fallback_doc = _doc("missing") + failed_fetcher = _FailedFetcher({"missing": fallback_doc}) + loader.failed_fetcher = failed_fetcher + + async def _fetch(source): # ruff:ignore[unused-async] + if source == "missing": + return None + return _doc(source) + + async def _fetch_html_docs(docs): # ruff:ignore[unused-async] + return docs + + monkeypatch.setattr(loader, "fetch", _fetch) + monkeypatch.setattr( + loader, + "_fetch_html_docs_again_using_playwright", + _fetch_html_docs, + ) + + docs = await loader.fetch_all("kept", "missing") + + assert [doc.attrs["source"] for doc in docs] == ["kept", "missing"] + assert failed_fetcher.calls == [("missing",)] + + +@pytest.mark.asyncio +async def test_docling_web_loader_passes_configured_deadline(monkeypatch): + """Configured Docling deadlines should reach the converter""" + captured = {} + loader = AsyncDoclingWebFileLoader( + pdf_pipeline_options={"document_timeout": 120} + ) + loader.content_fetcher = _Fetcher() + + async def _read_docling_web_file(*args, **kwargs): + await asyncio.sleep(0) + captured.update(kwargs) + return _doc("https://example.com/sample.pdf") + + monkeypatch.setattr( + "compass.web.file_loader.read_docling_web_file", + _read_docling_web_file, + ) + + doc, raw_content = await loader._fetch_doc( + "https://example.com/sample.pdf" + ) + + assert doc.attrs["doc_type"] == "pdf" + assert raw_content == b"content" + assert captured["pdf_pipeline_options"] == {"document_timeout": 120} + + +if __name__ == "__main__": + pytest.main(["-q", "--show-capture=all", Path(__file__), "-rapP"])