src/transform/extract.py

import random

import os
import sys
import time
import queue
import logging
import gc
import multiprocessing as mp
import argparse
from pathlib import Path
import torch

def configure_parallelism(workers_per_gpu=5):
    """
    Tuned specifically for 96GB VRAM GPUs.
    We cap workers to prevent 'Thundering Herd' and use VRAM for 'Batch Power'.
    """
    if not torch.cuda.is_available():
        return max(1, mp.cpu_count() // 2), 1, 0

    num_gpus = torch.cuda.device_count()
    gpu_properties = torch.cuda.get_device_properties(0)
    total_vram_gb = gpu_properties.total_memory / (1024**3)
    
    # --- THE STABILITY STRATEGY ---
    # On 96GB, 8-10 workers is the "Sweet Spot". 
    # More workers than this creates too much 'context switching' overhead on the GPU. 
    
    total_slots = num_gpus * workers_per_gpu
    
    print(f"🔍 GH200/A100 Detected: {num_gpus} GPUs | {total_vram_gb:.1f} GB VRAM")
    print(f"⚙️  Stability Config: {workers_per_gpu} workers/GPU | {total_slots} Total Slots")
    
    # --- SYSTEM TUNING ---
    os.environ["OLLAMA_NUM_PARALLEL"] = str(total_slots)
    os.environ["OLLAMA_MAX_QUEUE"] = "2048" # Large buffer for requests
    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    
    return total_slots, workers_per_gpu, num_gpus



# --- CRITICAL: MULTIPROCESSING SETUP ---
try:
    mp.set_start_method('spawn', force=True)
except RuntimeError:
    pass 

from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
from marker.config.parser import ConfigParser

# Configure Logger
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - [GPU-%(processName)s] - %(message)s',
    datefmt='%H:%M:%S'
)
logger = logging.getLogger(__name__)

def worker_routine(worker_id, gpu_id, batch_queue, output_dir, ollama_config, marker_config):
    """
    Optimized Worker for GH200:
    1. Receives a BATCH of files (List) to reduce queue overhead.
    2. Uses torch.compile for architectural optimization.
    3. Skips image extraction for speed.
    4. Fails fast on tables (Triage) to prevent LLM stalls.
    """
    
    time.sleep(worker_id * 1.5) 
    
    mp.current_process().name = f"{worker_id}:Dev{gpu_id}"
    logger.info(f"Initializing Worker {worker_id}...")
    
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    os.environ["TORCH_DEVICE"] = f"cuda:{gpu_id}"
    
    # 2. Model Initialization
    try:
        # Load Surya/OCR weights
        artifact_dict = create_model_dict(
            device=f"cuda:{gpu_id}", 
            dtype=torch.bfloat16, 
            attention_implementation="flash_attention_2"
        )

        # --- OPTIMIZATION 1: torch.compile ---
        logger.info("Compiling models with torch.compile... (One-time setup)")
        for key, model in artifact_dict.items():
            if hasattr(model, 'forward'):
                artifact_dict[key] = torch.compile(model, mode="max-autotune")

        # --- OPTIMIZATION 2: Config Tuning ---
        full_config = {
            "output_format": "markdown",
            "disable_multiprocessing": True,
            "extract_images": False,         # Speed up: Skip image extraction
            "ocr_all_pages": False,
            "use_llm": True,
            "llm_service": "marker.services.ollama.OllamaService",
            
            # --- TRIAGE STRATEGY ---
            "max_table_retries": 0,          # Fail fast if table extraction stalls
            "llm_service_timeout": 150, # Don't let a table hold a worker for more than x minutes
            
            **ollama_config,
            **marker_config
        }
        
        config_parser = ConfigParser(full_config)
        
        converter = PdfConverter(
            config=config_parser.generate_config_dict(),
            artifact_dict=artifact_dict,
            processor_list=config_parser.get_processors(),
            renderer=config_parser.get_renderer(),
            llm_service=config_parser.get_llm_service()
        )
        logger.info(f"Worker Ready. Waiting for batches...")
        
    except Exception as e:
        logger.error(f"Initialization Failed: {e}")
        return

    # 3. Batch Work Loop
    batches_processed = 0
    while True:
        try:
            # Get a list of files (Batch)
            batch_files = batch_queue.get(timeout=5)
        except queue.Empty:
            logger.info(f"Queue empty. Worker shutting down. Processed {batches_processed} batches.")
            break
            
        # Process the batch locally
        for pdf_path_str in batch_files:
            try:
                pdf_path = Path(pdf_path_str)
                doc_out_dir = Path(output_dir) / pdf_path.stem
                md_file = doc_out_dir / f"{pdf_path.stem}.md"

                if md_file.exists():
                    continue

                # Heavy Compute (OCR + LLM)
                rendered = converter(str(pdf_path))
                
                if rendered is None:
                    logger.warning(f"Skipping {pdf_path.name}: Converter returned None")
                    continue

                text, meta, images = text_from_rendered(rendered)

                # Write Output
                doc_out_dir.mkdir(parents=True, exist_ok=True)
                with open(md_file, "w", encoding="utf-8") as f:
                    f.write(text)
            
            except Exception as e:
                logger.error(f"Failed on {pdf_path.name}: {e}")
        
        # Cleanup after batch to keep VRAM healthy
        batches_processed += 1
        
        # Aggressive GC after every batch prevents "memory creep" on long runs
        gc.collect()
        torch.cuda.empty_cache()

class MarkerFolderProcessor:
    def __init__(self, output_dir, ollama_url, ollama_model, batch_multiplier, workers_per_gpu, num_gpus):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.num_gpus = num_gpus
        
        # We now accept the dynamic num_gpus passed from __main__
        if self.num_gpus > 0:
            print(f" Detected {self.num_gpus} GPUs (Dynamic Mode)")
        else:
            print(" No GPUs detected. Running in CPU mode.")

        self.workers_per_gpu = workers_per_gpu
        
        # Configs passed to workers
        self.ollama_config = {
            "ollama_base_url": ollama_url,
            "ollama_model": ollama_model,
            "ollama_timeout": 600, # 10 mins max per request
            "ollama_options": {
                "num_ctx": 32768,
                "num_predict": 2048,
                "temperature": 0.0
            }
        }
        self.marker_config = {"batch_multiplier": batch_multiplier}

    def process_folder(self, source_folder, batch_size=10, subset=None):
        if subset is not None:
            # Use the partitioned list provided by run_transform.py
            pdfs = [Path(p) for p in subset]
        else:
            source_path = Path(source_folder)
            pdfs = sorted(list(source_path.glob("*.pdf")))
        
        if not pdfs:
            print("No PDFs to process.")
            return

        manager = mp.Manager()
        batch_queue = manager.Queue()
        
        # --- BATCHING STRATEGY ---
        # Chunk the list of PDFs into batches
        chunks = [pdfs[i:i + batch_size] for i in range(0, len(pdfs), batch_size)]
        print(f" Created {len(chunks)} batches of {batch_size} files each.")
        
        for chunk in chunks:
            batch_queue.put([str(p) for p in chunk])

        total_workers = (self.num_gpus * self.workers_per_gpu) if self.num_gpus > 0 else 1
        print(f" Launching {total_workers} workers on {self.num_gpus} GPUs...")

        processes = []
        for i in range(total_workers):
            gpu_id = i % self.num_gpus if self.num_gpus > 0 else 0
            p = mp.Process(
                target=worker_routine,
                args=(i, gpu_id, batch_queue, self.output_dir, self.ollama_config, self.marker_config)
            )
            p.start()
            processes.append(p)

        for p in processes:
            p.join()
            
        print(" Extraction Complete.")

# This block only runs if you execute 'python extract.py' directly.
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="Input folder of PDFs")
    parser.add_argument("--output", required=True, help="Output folder")
    parser.add_argument("--url", default="http://localhost:11434", help="Ollama URL")
    parser.add_argument("--model", default="llama3", help="Ollama Model Name")
    args = parser.parse_args()

    # --- DYNAMIC HARDWARE DETECTION ---
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        gpu_properties = torch.cuda.get_device_properties(0)
        total_vram_gb = gpu_properties.total_memory / (1024**3)
        
        # Calculate optimal workers: (VRAM - 2GB overhead) / 5GB per worker
        workers_per_gpu = int((total_vram_gb - 2) // 5)
        workers_per_gpu = max(1, workers_per_gpu) # Minimum 1
        
        total_slots = num_gpus * workers_per_gpu
        print(f"⚙️  Dynamic Config: {num_gpus} GPUs | {workers_per_gpu} workers/GPU | {total_slots} Total Slots")
        
        # Set Env vars for external tools (optional, but good practice)
        os.environ["OLLAMA_NUM_PARALLEL"] = str(total_slots)
        
    else:
        num_gpus = 0
        workers_per_gpu = 1
        print("  No GPU detected. Defaulting to 1 worker.")

    # --- PROCESSOR INIT ---
    processor = MarkerFolderProcessor(
        output_dir=args.output,
        ollama_url=args.url,
        ollama_model=args.model,
        batch_multiplier=2,      
        workers_per_gpu=workers_per_gpu, # Passed dynamically
        num_gpus=num_gpus                # Passed dynamically
    )
    
    processor.process_folder(args.input, batch_size=10)