diff --git a/rubbernecker/parse/tool.py b/rubbernecker/parse/tool.py index 1834509..504f35b 100644 --- a/rubbernecker/parse/tool.py +++ b/rubbernecker/parse/tool.py @@ -4,13 +4,18 @@ from .base import Parser, list_parsers from avrokit import URL, parse_url, create_url_mapping, avro_reader, avro_writer +from avro.schema import Schema from dataclasses import dataclass import argparse import importlib import importlib.util import logging +import multiprocessing +import os import sys +import time from pathlib import Path +from typing import Any, cast logger = logging.getLogger("parsetool") @@ -21,14 +26,114 @@ class ParseToolStats: count_output: int = 0 count_error: int = 0 + def __add__(self, other: "ParseToolStats") -> "ParseToolStats": + return ParseToolStats( + count_input=self.count_input + other.count_input, + count_output=self.count_output + other.count_output, + count_error=self.count_error + other.count_error, + ) + + +POISON_PILL = (None, None) # Signal to workers to stop +WRITE_DONE = (None, None, None) # Signal that writing is complete + + +def worker_process( + parser_class: str, + script_path: str | None, + work_queue: multiprocessing.Queue, + result_queue: multiprocessing.Queue, +) -> None: + """Worker process that parses records from work_queue and puts results on result_queue.""" + logger.debug("Worker starting") + try: + # Initialize parser in worker process + parser = ParseTool._load_parser_static(parser_class, script_path) + logger.debug("Worker initialized parser successfully") + except Exception as e: + logger.error("Worker failed to initialize parser: %s", e) + # Signal done so writer doesn't hang + result_queue.put(("done", None, None, None)) + return + + while True: + task = work_queue.get() + if task == POISON_PILL: + # Signal to writer that this worker is done + result_queue.put(("done", None, None, None)) + break + + seq_id, record = task + stats = ParseToolStats() + + try: + stats.count_input = 1 + results: list[dict[str, Any]] = [] + for parsed_record in parser.parse(record): + if parsed_record is not None: + results.append(cast(dict[str, Any], parsed_record)) + stats.count_output += 1 + result_queue.put(("result", seq_id, results, stats)) + except Exception as e: + logger.error("Error parsing record: %s", e) + if logging.DEBUG == logger.getEffectiveLevel(): + logger.exception(e) + stats.count_error = 1 + result_queue.put(("result", seq_id, None, stats)) + + +def writer_process( + output_url: URL, + schema: Schema, + result_queue: multiprocessing.Queue, + num_workers: int, + stats_accumulator: dict, +) -> None: + """Writer process that consumes results and writes to Avro.""" + logger.debug("Writer process starting, expecting %d workers", num_workers) + workers_done = 0 + + with avro_writer(output_url.with_mode("wb"), schema) as writer: + while workers_done < num_workers: + msg_type, seq_id, results, task_stats = result_queue.get() + + if msg_type == "done": + workers_done += 1 + continue + + # Accumulate stats + stats_accumulator["count_input"] += task_stats.count_input + stats_accumulator["count_output"] += task_stats.count_output + stats_accumulator["count_error"] += task_stats.count_error + + # Write results + if results: + for parsed_record in results: + writer.append(parsed_record) + + # Progress logging every 1k records + total_processed = ( + stats_accumulator["count_output"] + stats_accumulator["count_error"] + ) + if total_processed % 1000 == 0: + logger.info( + "Progress: input=%d output=%d errors=%d", + stats_accumulator["count_input"], + stats_accumulator["count_output"], + stats_accumulator["count_error"], + ) + + logger.debug("Writer process finished") + class ParseTool: def name(self) -> str: return "parse" - def load_parser(self, class_name: str, script_path: str | None = None) -> Parser: + @staticmethod + def _load_parser_static(class_name: str, script_path: str | None = None) -> Parser: + """Static version of load_parser for use in worker processes.""" if script_path: - # Load module from arbitrary Python file script_file = Path(script_path).resolve() if not script_file.exists(): raise FileNotFoundError(f"Script file not found: {script_file}") @@ -41,11 +146,13 @@ def load_parser(self, class_name: str, script_path: str | None = None) -> Parser spec.loader.exec_module(module) return getattr(module, class_name)() else: - # Load from installed package module_name, _, class_name = class_name.rpartition(".") module = importlib.import_module(module_name) return getattr(module, class_name)() + def load_parser(self, class_name: str, script_path: str | None = None) -> Parser: + return self._load_parser_static(class_name, script_path) + def list_parsers(self) -> list[str]: acc: list[str] = [] for parser in list_parsers(): @@ -87,6 +194,139 @@ def parse( logger.debug("%s", stats) return stats + def parse_parallel( + self, + parser: Parser, + base_input_url: URL, + base_output_url: URL, + num_workers: int, + script_path: str | None = None, + ) -> ParseToolStats: + """Parse with process-level parallelism using bounded queues.""" + # Get parser info for worker initialization + parser_class = parser.__class__.__name__ + # When using a script, just use the class name (module is derived from script filename) + # When using built-in parser, use full module.class path + if script_path: + worker_parser_name = parser_class + else: + parser_module = parser.__class__.__module__ + worker_parser_name = f"{parser_module}.{parser_class}" + + stats = ParseToolStats() + + for input_url, output_url in create_url_mapping( + base_input_url, base_output_url + ): + logger.info( + "Parsing (parallel=%d) %s -> %s", num_workers, input_url, output_url + ) + + # Bounded queues for backpressure (larger to prevent deadlock) + queue_maxsize = max(100, num_workers * 10) + work_queue: multiprocessing.Queue = multiprocessing.Queue( + maxsize=queue_maxsize + ) + result_queue: multiprocessing.Queue = multiprocessing.Queue( + maxsize=queue_maxsize + ) + + # Shared stats accumulator + manager = multiprocessing.Manager() + shared_stats = manager.dict() + shared_stats["count_input"] = 0 + shared_stats["count_output"] = 0 + shared_stats["count_error"] = 0 + + # Start worker processes + workers: list[multiprocessing.Process] = [] + for i in range(num_workers): + p = multiprocessing.Process( + target=worker_process, + args=(worker_parser_name, script_path, work_queue, result_queue), + ) + p.start() + workers.append(p) + logger.debug("Started worker %d (pid=%d)", i, p.pid) + + # Give workers a moment to initialize + time.sleep(0.5) + + # Check if any workers died immediately + dead_workers = [i for i, p in enumerate(workers) if not p.is_alive()] + if dead_workers: + raise RuntimeError( + f"Workers {dead_workers} died immediately after starting" + ) + + # Start writer process + writer_p = multiprocessing.Process( + target=writer_process, + args=( + output_url, + parser.schema(), + result_queue, + num_workers, + shared_stats, + ), + ) + writer_p.start() + + try: + # Producer: read records and put on work queue (blocks when full) + record_count = 0 + logger.info("Starting to read records from %s", input_url) + with avro_reader(input_url.with_mode("rb")) as reader: + logger.info("Avro reader opened successfully") + for seq_id, record in enumerate(reader): + work_queue.put((seq_id, record)) + record_count += 1 + + logger.info("Finished queuing %d records", record_count) + + # Signal workers to stop + for _ in range(num_workers): + work_queue.put(POISON_PILL) + + # Wait for workers to finish (no timeout - let them work) + logger.info("Waiting for %d workers to finish...", num_workers) + for i, p in enumerate(workers): + p.join() + logger.debug("Worker %d finished", i) + + # Wait for writer to finish (no timeout) + logger.info("Waiting for writer to finish...") + writer_p.join() + logger.debug("Writer finished") + + # Copy stats from shared dict + stats.count_input += shared_stats["count_input"] + stats.count_output += shared_stats["count_output"] + stats.count_error += shared_stats["count_error"] + + except KeyboardInterrupt: + logger.warning("Interrupted! Shutting down workers...") + + # Terminate workers + for p in workers: + if p.is_alive(): + p.terminate() + p.join(timeout=1) + + if writer_p.is_alive(): + writer_p.terminate() + writer_p.join(timeout=1) + + # Clean up partial output + logger.warning("Discarding incomplete output for %s", output_url) + try: + output_url.delete() + except Exception: + pass + raise + + return stats + def configure(self, subparsers: argparse._SubParsersAction) -> None: parser = subparsers.add_parser(self.name(), help="Parse output of crawl") parser.add_argument( @@ -106,6 +346,17 @@ def configure(self, subparsers: argparse._SubParsersAction) -> None: help="Path to a Python file containing the parser implementation", ) + # Calculate default parallelism + cpu_count = os.cpu_count() + default_parallelism = max(1, (cpu_count // 2) if cpu_count else 1) + + parser.add_argument( + "--parallelism", + type=int, + default=default_parallelism, + help=f"Number of parallel worker processes (default: {default_parallelism})", + ) + def run(self, args: argparse.Namespace) -> None: # Validate built-in parser name if no script provided if not args.script and args.name not in self.list_parsers(): @@ -113,9 +364,16 @@ def run(self, args: argparse.Namespace) -> None: f"Unknown parser: {args.name}. " f"Use --script to load from a file, or choose from: {', '.join(self.list_parsers())}" ) - stats = self.parse( - self.load_parser(args.name, args.script), - parse_url(args.input_url), - parse_url(args.output_url), - ) + + loaded_parser = self.load_parser(args.name, args.script) + input_url = parse_url(args.input_url) + output_url = parse_url(args.output_url) + + if args.parallelism > 1: + stats = self.parse_parallel( + loaded_parser, input_url, output_url, args.parallelism, args.script + ) + else: + stats = self.parse(loaded_parser, input_url, output_url) + logger.info("%s (done)", stats)