examples/batch_lightrag_ood_context_generation.py

import os
import json
import time
import asyncio
import logging
import logging.config
import multiprocessing as mp
from pathlib import Path
from functools import partial

from lightrag import LightRAG, QueryParam
from lightrag.llm.hf import hf_model_complete
from lightrag.llm.openai import openai_embed
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.utils import logger, set_verbose_debug
from lightrag.rerank import ali_rerank

WORKING_DIR = "/root/githubs/LightRAG/ood_novel_"
QA_PAIR_FILE = "/root/githubs/LightRAG/generated_data/batch_neg_ood_qa_pair.json"
OUTPUT_FILE = "/root/githubs/LightRAG/generated_data/batch_neg_ood_qa_pair_context_single.json"
DEFAULT_GPU_IDS = ["7"]
OPENAI_API_KEY = "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"


def configure_logging():
    """Configure logging for the application"""

    # Reset any existing handlers to ensure clean configuration
    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
        logger_instance = logging.getLogger(logger_name)
        logger_instance.handlers = []
        logger_instance.filters = []

    # Get log directory path from environment variable or use current directory
    log_dir = os.getenv("LOG_DIR", os.getcwd())
    log_file_path = os.path.abspath(os.path.join(log_dir, "lightrag_demo.log"))

    print(f"\nLightRAG demo log file: {log_file_path}\n")
    os.makedirs(os.path.dirname(log_dir), exist_ok=True)

    # Get log file max size and backup count from environment variables
    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups

    logging.config.dictConfig(
        {
            "version": 1,
            "disable_existing_loggers": False,
            "formatters": {
                "default": {
                    "format": "%(levelname)s: %(message)s",
                },
                "detailed": {
                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                },
            },
            "handlers": {
                "console": {
                    "formatter": "default",
                    "class": "logging.StreamHandler",
                    "stream": "ext://sys.stderr",
                },
                "file": {
                    "formatter": "detailed",
                    "class": "logging.handlers.RotatingFileHandler",
                    "filename": log_file_path,
                    "maxBytes": log_max_bytes,
                    "backupCount": log_backup_count,
                    "encoding": "utf-8",
                },
            },
            "loggers": {
                "lightrag": {
                    "handlers": ["console", "file"],
                    "level": "INFO",
                    "propagate": False,
                },
            },
        }
    )

    # Set the logger level to INFO
    logger.setLevel(logging.INFO)
    # Enable verbose debug if needed
    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")


if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)


rerank_model_func = partial(
    ali_rerank,
    model="gte-rerank-v2",
    api_key="sk-378f9fd16e3148769b992d45a94001ad",
)


async def initialize_rag():
    rag = LightRAG(
        working_dir=WORKING_DIR,
        llm_model_func=hf_model_complete,
        llm_model_name="Qwen/Qwen3-4B-Instruct-2507",
        embedding_func=openai_embed,
        rerank_model_func=rerank_model_func,
    )

    await rag.initialize_storages()
    await initialize_pipeline_status()

    return rag


def parse_gpu_ids(value):
    """Parse a comma separated GPU string into a clean list."""
    return [gpu.strip() for gpu in value.split(",") if gpu.strip()]


def determine_gpu_ids(default_gpu_ids=None):
    """Resolve which GPU ids should be used for processing."""
    default_gpu_ids = default_gpu_ids or DEFAULT_GPU_IDS
    requested_visible = os.getenv("QWEN3_VISIBLE_GPUS")
    if requested_visible:
        gpu_ids = parse_gpu_ids(requested_visible)
        if gpu_ids:
            print(f"Using GPUs from QWEN3_VISIBLE_GPUS={','.join(gpu_ids)}")
        else:
            gpu_ids = list(default_gpu_ids)
            print(
                f"QWEN3_VISIBLE_GPUS was empty; defaulting to GPUs {','.join(gpu_ids)}."
            )
    else:
        gpu_ids = list(default_gpu_ids)
        print(f"QWEN3_VISIBLE_GPUS not set; defaulting to GPUs {','.join(gpu_ids)}.")

    visible_str = ",".join(gpu_ids)
    os.environ["CUDA_VISIBLE_DEVICES"] = visible_str
    os.environ["QWEN3_VISIBLE_GPUS"] = visible_str
    return gpu_ids


def split_qa_pairs(qa_pairs, num_chunks):
    """Split QA pairs into roughly even chunks for multi-GPU processing."""
    if not qa_pairs:
        return []

    num_chunks = max(1, num_chunks)
    total = len(qa_pairs)
    base, extra = divmod(total, num_chunks)
    chunks = []
    start = 0
    for idx in range(num_chunks):
        stop = start + base + (1 if idx < extra else 0)
        if start >= total:
            break
        chunks.append(qa_pairs[start:stop])
        start = stop

    return chunks


def load_qa_pairs(qa_pair_file):
    """Load QA pairs from disk and enrich them with metadata."""
    with open(qa_pair_file, "r", encoding="utf-8") as fh:
        qa_pair_dict = json.load(fh)

    qa_pairs = []
    for char_name, pairs in qa_pair_dict.items():
        for qa_pair in pairs:
            enriched = dict(qa_pair)
            enriched["charactor_name"] = char_name
            enriched["_index"] = len(qa_pairs)
            qa_pairs.append(enriched)

    return qa_pairs


async def process_qa_pairs(gpu_id, qa_pairs):
    """Execute LightRAG queries for a chunk of QA pairs on a single GPU."""
    rag = await initialize_rag()
    processed = []
    mode = "hybrid_context"
    start_time = time.time()

    for qa_pair in qa_pairs:
        try:
            search_results, ll_keywords, hl_keywords = await rag.aquery(
                qa_pair["question"],
                param=QueryParam(mode=mode),
                timeline_key=qa_pair.get("timeline_key"),
                charactor_name=qa_pair.get("charactor_name"),
            )
            qa_pair["final_entities"] = search_results["final_entities"]
            qa_pair["final_relations"] = search_results["final_relations"]
            qa_pair["hl_keywords"] = hl_keywords
            qa_pair["ll_keywords"] = ll_keywords
        except Exception as exc:
            qa_pair["error"] = str(exc)
            logger.exception(
                "[GPU %s] Failed to process question: %s",
                gpu_id,
                qa_pair.get("question", "unknown"),
            )

        processed.append(qa_pair)

    elapsed = time.time() - start_time
    logger.info(
        "[GPU %s] Processed %d questions in %.2f seconds.",
        gpu_id,
        len(qa_pairs),
        elapsed,
    )
    return processed


def worker_process(gpu_id, qa_pairs, result_queue, base_log_dir):
    """Worker entry point for multi-process execution."""
    if not qa_pairs:
        result_queue.put({"gpu": gpu_id, "results": [], "error": None})
        return

    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
    os.environ["HF_VISIBLE_GPUS"] = gpu_id
    os.environ["QWEN3_VISIBLE_GPUS"] = gpu_id

    process_log_dir = Path(base_log_dir) / f"gpu_{gpu_id}"
    process_log_dir.mkdir(parents=True, exist_ok=True)
    os.environ["LOG_DIR"] = str(process_log_dir)

    configure_logging()
    print(f"[GPU {gpu_id}] Worker handling {len(qa_pairs)} questions.")

    try:
        processed = asyncio.run(process_qa_pairs(gpu_id, qa_pairs))
        result_queue.put({"gpu": gpu_id, "results": processed, "error": None})
    except Exception as exc:
        logger.exception("Worker %s encountered an unrecoverable error.", gpu_id)
        result_queue.put({"gpu": gpu_id, "results": [], "error": str(exc)})


def run_multiprocess(qa_pairs, gpu_ids, log_dir):
    """Distribute QA pairs across GPUs and collect results."""
    chunks = split_qa_pairs(qa_pairs, len(gpu_ids))
    if not chunks:
        return [], []

    ctx = mp.get_context("spawn")
    result_queue = ctx.Queue()
    processes = []

    for gpu_id, chunk in zip(gpu_ids, chunks):
        if not chunk:
            continue
        proc = ctx.Process(
            target=worker_process,
            args=(gpu_id, chunk, result_queue, log_dir),
        )
        proc.start()
        processes.append(proc)

    collected = []
    errors = []

    for _ in processes:
        payload = result_queue.get()
        collected.extend(payload["results"])
        if payload["error"]:
            errors.append((payload["gpu"], payload["error"]))

    for proc in processes:
        proc.join()

    return collected, errors


def main():
    os.environ.setdefault("LOG_DIR", os.getcwd())
    try:
        mp.set_start_method("spawn")
    except RuntimeError:
        pass

    configure_logging()

    gpu_ids = determine_gpu_ids()
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

    qa_pairs = load_qa_pairs(QA_PAIR_FILE)
    if not qa_pairs:
        print("No QA pairs found to process.")
        return

    base_log_dir = os.getenv("LOG_DIR", os.getcwd())
    start_time = time.time()
    processed_pairs, worker_errors = run_multiprocess(
        qa_pairs, gpu_ids, base_log_dir
    )

    if len(processed_pairs) != len(qa_pairs):
        print(
            f"Warning: expected {len(qa_pairs)} results but received {len(processed_pairs)}."
        )

    processed_pairs.sort(key=lambda item: item["_index"])
    for pair in processed_pairs:
        pair.pop("_index", None)

    Path(OUTPUT_FILE).write_text(
        json.dumps(processed_pairs, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    elapsed = time.time() - start_time
    print(f"Total time taken for processing: {elapsed} seconds")
    if worker_errors:
        print("Some workers reported errors:")
        for gpu_id, err in worker_errors:
            print(f" - GPU {gpu_id}: {err}")


if __name__ == "__main__":
    main()
    print("\nDone!")