RunPod embed v3 - fixed cleanup

runpod_embed_v3.py

"""
RunPod A100 Batch Embedding Script — v3 (File-Based Pipeline)
Eliminates bore tunnel bottleneck by decoupling DB I/O from GPU.

STRATEGY: Pull → Embed → Push (in chunks)
1. PULL: Stream N docs from remote DB to local CSV (one-time bore cost)
2. EMBED: Read local CSV, encode on GPU at full speed (no network!)
3. PUSH: Bulk-write embeddings back to remote DB (one-time bore cost)
4. Repeat for next chunk

Usage:
    pip install sentence-transformers psycopg2-binary
    python runpod_embed_v3.py --db-host bore.pub --db-port 16870 --chunk-size 2000000
"""

import argparse
import csv
import gc
import gzip
import json
import logging
import os
import time
import sys

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s: %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('/workspace/embed.log'),
    ]
)
logger = logging.getLogger(__name__)

PROGRESS_FILE = '/workspace/embed_progress.json'
WORK_DIR = '/workspace/chunks'


def save_progress(state):
    with open(PROGRESS_FILE, 'w') as f:
        json.dump(state, f, indent=2)


def load_progress():
    if os.path.exists(PROGRESS_FILE):
        with open(PROGRESS_FILE) as f:
            return json.load(f)
    return {'chunk': 0, 'total_done': 0, 'total_docs': 0, 'errors': 0, 'start_time': time.time()}


def get_db_connection(args):
    import psycopg2
    for attempt in range(15):
        try:
            conn = psycopg2.connect(
                host=args.db_host, port=args.db_port,
                dbname=args.db_name, user=args.db_user,
                password=args.db_pass, connect_timeout=15
            )
            conn.autocommit = False
            return conn
        except Exception as e:
            delay = min(5 * (2 ** min(attempt, 4)), 60)
            logger.warning("DB connect failed (attempt %d): %s. Retry in %ds", attempt+1, e, delay)
            time.sleep(delay)
    raise RuntimeError("Cannot connect to DB after 15 attempts")


def phase_pull(conn, chunk_id, offset, chunk_size, max_chars):
    """PULL: Stream docs from remote DB to local CSV file."""
    os.makedirs(WORK_DIR, exist_ok=True)
    csv_path = os.path.join(WORK_DIR, f'docs_{chunk_id:04d}.csv.gz')

    if os.path.exists(csv_path):
        # Count rows in existing file
        count = 0
        with gzip.open(csv_path, 'rt') as f:
            reader = csv.reader(f)
            for _ in reader:
                count += 1
        if count > 0:
            logger.info("PULL: Chunk %d already exists (%d rows), skipping", chunk_id, count)
            return csv_path, count

    logger.info("PULL: Streaming chunk %d (offset=%d, limit=%d) from DB...", chunk_id, offset, chunk_size)
    t0 = time.time()

    import io

    # Use COPY TO STDOUT for maximum streaming throughput.
    # COPY uses PostgreSQL's binary protocol — 5-10x faster than SELECT+fetchmany
    # for large result sets through a network tunnel.
    # S51 FIX: Use keyset pagination (WHERE id > last_id) instead of OFFSET.
    # OFFSET forces PostgreSQL to scan and skip N rows — O(N) per chunk.
    # Keyset pagination uses the primary key index — O(log N) per chunk.
    # For chunk 10 with OFFSET, PG skips 10M rows. With keyset, it seeks directly.
    last_id_file = os.path.join(WORK_DIR, 'last_seen_id.txt')
    last_id = None
    if offset > 0 and os.path.exists(last_id_file):
        with open(last_id_file) as f:
            last_id = f.read().strip()

    try:
        if last_id:
            # Keyset pagination — fast seek via primary key index
            copy_query = (
                f"COPY (SELECT id, left(content, {max_chars}) FROM documents "
                f"WHERE embedding_dense IS NULL AND id > '{last_id}' ORDER BY id "
                f"LIMIT {chunk_size}) TO STDOUT CSV"
            )
        else:
            # First chunk — no last_id yet, use simple LIMIT
            copy_query = (
                f"COPY (SELECT id, left(content, {max_chars}) FROM documents "
                f"WHERE embedding_dense IS NULL ORDER BY id "
                f"LIMIT {chunk_size}) TO STDOUT CSV"
            )
        count = 0
        with gzip.open(csv_path, 'wt', compresslevel=3) as f:
            cur = conn.cursor()
            cur.copy_expert(copy_query, f)
            conn.commit()
        # Count rows and save last ID for next chunk's keyset pagination
        last_row_id = None
        with gzip.open(csv_path, 'rt') as f:
            for line in f:
                count += 1
                parts = line.strip().split(',', 1)
                if parts:
                    last_row_id = parts[0]
        if last_row_id:
            with open(last_id_file, 'w') as f:
                f.write(last_row_id)
        logger.info("PULL: Used COPY TO STDOUT (%d rows, keyset=%s)", count, bool(last_id))
    except Exception as copy_err:
        logger.warning("PULL: COPY failed (%s), falling back to cursor fetch", copy_err)
        # Fallback: server-side cursor with large fetch
        cur = conn.cursor('pull_cursor')
        cur.itersize = 50000
        cur.execute("""
            SELECT id, left(content, %s) FROM documents
            WHERE embedding_dense IS NULL
            ORDER BY id
            LIMIT %s OFFSET %s
        """, (max_chars, chunk_size, offset))
        count = 0
        with gzip.open(csv_path, 'wt', compresslevel=3) as f:
            writer = csv.writer(f)
            for row in cur:
                doc_id, content = row
                text = (content or "").strip()
                if text:
                    writer.writerow([doc_id, text])
                    count += 1
        cur.close()
    elapsed = time.time() - t0
    size_mb = os.path.getsize(csv_path) / 1024 / 1024
    logger.info("PULL: Chunk %d done — %d rows, %.1f MB, %.1fs (%.0f rows/s)",
                chunk_id, count, size_mb, elapsed, count/elapsed if elapsed > 0 else 0)
    return csv_path, count


def phase_embed(csv_path, model, chunk_id, gpu_batch):
    """EMBED: Read local CSV, encode on GPU, write embeddings to local file."""
    emb_path = os.path.join(WORK_DIR, f'emb_{chunk_id:04d}.csv.gz')

    if os.path.exists(emb_path):
        count = 0
        with gzip.open(emb_path, 'rt') as f:
            reader = csv.reader(f)
            for _ in reader:
                count += 1
        if count > 0:
            logger.info("EMBED: Chunk %d embeddings already exist (%d rows), skipping", chunk_id, count)
            return emb_path, count

    logger.info("EMBED: Loading chunk %d from %s...", chunk_id, csv_path)
    t0 = time.time()

    # Read all docs from CSV
    doc_ids = []
    texts = []
    with gzip.open(csv_path, 'rt') as f:
        reader = csv.reader(f)
        for row in reader:
            doc_ids.append(row[0])
            texts.append(row[1])

    load_time = time.time() - t0
    logger.info("EMBED: Loaded %d texts in %.1fs. Encoding on GPU...", len(texts), load_time)

    # Encode on GPU in sub-batches
    t1 = time.time()
    import numpy as np
    embeddings = model.encode(
        texts,
        batch_size=gpu_batch,
        show_progress_bar=True,
        normalize_embeddings=True
    )
    encode_time = time.time() - t1
    rate = len(texts) / encode_time if encode_time > 0 else 0
    logger.info("EMBED: Encoded %d texts in %.1fs (%.0f/s)", len(texts), encode_time, rate)

    # Write embeddings to compressed CSV
    t2 = time.time()
    with gzip.open(emb_path, 'wt', compresslevel=3) as f:
        writer = csv.writer(f)
        for doc_id, emb in zip(doc_ids, embeddings):
            vec_str = "[" + ",".join(f"{float(v):.5f}" for v in emb) + "]"
            writer.writerow([doc_id, vec_str])

    write_time = time.time() - t2
    size_mb = os.path.getsize(emb_path) / 1024 / 1024
    logger.info("EMBED: Wrote embeddings to %s (%.1f MB) in %.1fs",
                emb_path, size_mb, write_time)

    # Free GPU memory
    del embeddings, texts
    gc.collect()
    import torch
    torch.cuda.empty_cache()

    return emb_path, len(doc_ids)


def phase_push(conn, emb_path, chunk_id):
    """PUSH: Bulk-write embeddings from local file back to remote DB.

    Uses COPY FROM into temp table + single bulk UPDATE join.
    This is 10-50x faster than execute_batch for large volumes because:
    - COPY streams data in binary protocol (minimal per-row overhead)
    - Single UPDATE...FROM join instead of N individual UPDATEs
    - Dramatically fewer network round-trips through bore tunnel
    """
    logger.info("PUSH: Writing chunk %d embeddings to DB (COPY+JOIN method)...", chunk_id)
    t0 = time.time()

    import io

    cur = conn.cursor()

    # Step 1: Create temp table (no indexes, minimal overhead)
    cur.execute("""
        CREATE TEMP TABLE IF NOT EXISTS _emb_staging (
            doc_id TEXT,
            vec_text TEXT
        ) ON COMMIT DROP
    """)
    cur.execute("TRUNCATE _emb_staging")

    # Step 2: Stream embeddings via COPY FROM STDIN (bulk binary protocol)
    t1 = time.time()
    copy_buf = io.StringIO()
    row_count = 0
    with gzip.open(emb_path, 'rt') as f:
        reader = csv.reader(f)
        for row in reader:
            doc_id = row[0]
            vec_str = row[1]
            # COPY CSV format: tab-separated, no quoting needed for numbers
            copy_buf.write(f"{doc_id}\t{vec_str}\n")
            row_count += 1

    copy_buf.seek(0)
    cur.copy_from(copy_buf, '_emb_staging', columns=('doc_id', 'vec_text'), sep='\t')
    copy_time = time.time() - t1
    logger.info("PUSH: COPY %d rows to staging in %.1fs (%.0f rows/s)",
                row_count, copy_time, row_count / copy_time if copy_time > 0 else 0)

    # Step 3: Single bulk UPDATE via JOIN (one query, not N queries)
    t2 = time.time()
    cur.execute("""
        UPDATE documents d
        SET embedding_dense = e.vec_text::vector
        FROM _emb_staging e
        WHERE d.id = e.doc_id::uuid
    """)
    updated = cur.rowcount
    conn.commit()
    update_time = time.time() - t2
    logger.info("PUSH: UPDATE %d rows in %.1fs (%.0f rows/s)",
                updated, update_time, updated / update_time if update_time > 0 else 0)

    elapsed = time.time() - t0
    rate = row_count / elapsed if elapsed > 0 else 0
    logger.info("PUSH: Chunk %d COMPLETE — %d rows in %.1fs total (%.0f/s, COPY=%.1fs, UPDATE=%.1fs)",
                chunk_id, row_count, elapsed, rate, copy_time, update_time)

    return row_count


def main():
    parser = argparse.ArgumentParser(description="GPU Batch Embedding v3 — File Pipeline")
    parser.add_argument('--db-host', default='bore.pub')
    parser.add_argument('--db-port', type=int, default=16870)
    parser.add_argument('--db-name', default='apra_unified')
    parser.add_argument('--db-user', default='apra')
    parser.add_argument('--db-pass', default='apra_dev')
    parser.add_argument('--chunk-size', type=int, default=2000000,
                       help='Documents per chunk (default: 2M)')
    parser.add_argument('--gpu-batch', type=int, default=256,
                       help='GPU encode sub-batch size (256 fits A100 in FP16)')
    parser.add_argument('--max-chars', type=int, default=1000,
                       help='Max chars per doc (1000 = ~250 tokens, good quality/speed balance)')
    parser.add_argument('--max-chunks', type=int, default=0,
                       help='Stop after N chunks (0=all)')
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--pull-only', action='store_true',
                       help='Only download docs, skip embed+push')
    parser.add_argument('--embed-only', action='store_true',
                       help='Only embed existing chunks, skip pull+push')
    parser.add_argument('--push-only', action='store_true',
                       help='Only push existing embeddings, skip pull+embed')
    args = parser.parse_args()

    from sentence_transformers import SentenceTransformer

    # Load model — FP16 for 2x encoding throughput on A100
    if not args.pull_only and not args.push_only:
        import torch
        logger.info("Loading BGE-M3 on GPU (FP16 for max throughput)...")
        t0 = time.time()
        model = SentenceTransformer(
            'BAAI/bge-m3', device='cuda',
            model_kwargs={'torch_dtype': torch.float16}
        )
        logger.info("Model loaded in %.1fs (dtype=%s)", time.time() - t0,
                     next(model.parameters()).dtype)
    else:
        model = None

    # Get total count
    logger.info("Connecting to DB for count...")
    conn = get_db_connection(args)
    cur = conn.cursor()
    cur.execute("SELECT count(*) FROM documents WHERE embedding_dense IS NULL")
    total = cur.fetchone()[0]
    logger.info("Total documents needing embedding: %d", total)

    # Resume state
    state = load_progress() if args.resume else {
        'chunk': 0, 'total_done': 0, 'total_docs': total,
        'errors': 0, 'start_time': time.time()
    }
    state['total_docs'] = total

    start_chunk = state.get('chunk', 0)
    total_done = state.get('total_done', 0)
    num_chunks = (total + args.chunk_size - 1) // args.chunk_size

    logger.info("Processing %d documents in %d chunks of %d (starting at chunk %d)",
                total, num_chunks, args.chunk_size, start_chunk)

    overall_start = time.time()

    import threading

    # S51 OPTIMIZATION: Overlap PUSH(N) with PULL(N+1) using threads.
    # PUSH writes to DB via bore tunnel while PULL reads next chunk.
    # GPU sits idle during both — this saves ~9 min per chunk.
    push_thread = None
    push_error = [None]  # Mutable container for thread error

    def threaded_push(push_conn, emb_path, cid, done_container, err_container):
        """Run PUSH in background thread so PULL can start immediately."""
        try:
            pushed = phase_push(push_conn, emb_path, cid)
            done_container[0] = pushed
        except Exception as e:
            err_container[0] = e
            logger.error("PUSH thread error: %s", e)

    for chunk_id in range(start_chunk, num_chunks):
        if args.max_chunks and chunk_id >= start_chunk + args.max_chunks:
            logger.info("Reached max_chunks=%d, stopping", args.max_chunks)
            break

        offset = chunk_id * args.chunk_size
        logger.info("=" * 60)
        logger.info("CHUNK %d/%d (offset=%d)", chunk_id + 1, num_chunks, offset)
        logger.info("=" * 60)

        try:
            # Wait for previous PUSH to finish before we reuse the connection
            if push_thread is not None:
                logger.info("Waiting for previous PUSH thread to complete...")
                push_thread.join()
                if push_error[0]:
                    logger.error("Previous PUSH failed: %s", push_error[0])
                push_thread = None

                # S51 FIX: Clean up previous chunk's temp files to prevent disk fill
                prev_chunk = chunk_id - 1
                for pattern in [f'docs_{prev_chunk:04d}.csv.gz', f'emb_{prev_chunk:04d}.csv.gz']:
                    old_file = os.path.join(WORK_DIR, pattern)
                    if os.path.exists(old_file):
                        os.remove(old_file)
                        logger.info("CLEANUP: Removed %s (%.1f MB freed)",
                                    pattern, os.path.getsize(old_file) / 1e6 if os.path.exists(old_file) else 0)

            # PHASE 1: PULL (runs while previous PUSH may still be finishing)
            if not args.embed_only and not args.push_only:
                try:
                    conn.cursor().execute("SELECT 1")
                except Exception:
                    conn = get_db_connection(args)
                csv_path, pulled = phase_pull(conn, chunk_id, offset, args.chunk_size, args.max_chars)
            else:
                csv_path = os.path.join(WORK_DIR, f'docs_{chunk_id:04d}.csv.gz')
                with gzip.open(csv_path, 'rt') as f:
                    pulled = sum(1 for _ in csv.reader(f))

            if pulled == 0:
                logger.info("No more documents, stopping")
                break

            # PHASE 2: EMBED (GPU — fast, no DB needed)
            if not args.pull_only and not args.push_only:
                emb_path, embedded = phase_embed(csv_path, model, chunk_id, args.gpu_batch)
            else:
                emb_path = os.path.join(WORK_DIR, f'emb_{chunk_id:04d}.csv.gz')

            # PHASE 3: PUSH — start in background thread so PULL(N+1) can begin
            if not args.pull_only and not args.embed_only:
                # Get a SEPARATE connection for the push thread
                push_conn = get_db_connection(args)
                push_done = [0]
                push_error = [None]
                push_thread = threading.Thread(
                    target=threaded_push,
                    args=(push_conn, emb_path, chunk_id, push_done, push_error),
                    daemon=True,
                )
                push_thread.start()
                logger.info("PUSH started in background thread — PULL(N+1) can begin immediately")
                total_done += pulled  # Optimistic — corrected if push fails
            else:
                total_done += pulled

            # DON'T clean up files yet — push thread still reading emb file
            # Cleanup happens at start of next iteration after push_thread.join()

            # Save progress
            elapsed = time.time() - overall_start
            rate = total_done / elapsed if elapsed > 0 else 0
            eta = (total - total_done) / rate / 3600 if rate > 0 else 0
            state.update({
                'chunk': chunk_id + 1,
                'total_done': total_done,
                'total_docs': total,
                'errors': state.get('errors', 0),
                'elapsed': elapsed,
                'rate': rate,
                'eta_hours': eta,
                'pct': total_done / total * 100 if total > 0 else 0,
                'updated': time.strftime('%Y-%m-%d %H:%M:%S'),
            })
            save_progress(state)

            logger.info("PROGRESS: %d/%d (%.2f%%) | %.0f/s overall | ETA: %.1fh | Chunk %d/%d",
                        total_done, total, total_done/total*100, rate, eta, chunk_id+1, num_chunks)

        except KeyboardInterrupt:
            logger.info("Interrupted. Progress saved.")
            state['chunk'] = chunk_id
            save_progress(state)
            break

        except Exception as e:
            logger.error("Chunk %d failed: %s", chunk_id, e, exc_info=True)
            state['errors'] = state.get('errors', 0) + 1
            state['chunk'] = chunk_id + 1  # Skip failed chunk
            save_progress(state)
            try:
                conn.close()
            except Exception:
                pass
            conn = get_db_connection(args)

    elapsed = time.time() - overall_start
    logger.info("=" * 60)
    logger.info("COMPLETE: %d/%d embedded in %.1fh (%.0f/s avg)",
                total_done, total, elapsed/3600, total_done/elapsed if elapsed > 0 else 0)
    logger.info("=" * 60)
    try:
        conn.close()
    except Exception:
        pass


if __name__ == "__main__":
    main()