app.py

import os
# 🔱 ANTI-DEADLOCK
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import json
import time
import asyncio
import unicodedata
import re
import string
import xxhash
import sqlite3
import pickle
import threading
import queue
import contextlib
import concurrent.futures
import zstandard as zstd
from collections import deque
from fastapi import FastAPI
from huggingface_hub import HfApi
from transformers import AutoTokenizer
from datasets import load_dataset
from datasketch import MinHash, MinHashLSH
from tenacity import retry, stop_after_attempt, wait_random_exponential
from pybloom_live import ScalableBloomFilter

# --- 🔱 CONFIGURATION ---
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("CRITICAL: HF_TOKEN environment variable is missing!")

PRIVATE_REPO = "Indro-ai/Indro-3B-Corpus"
DATASET_NAME = "HuggingFaceFW/fineweb-edu"
TOKENIZER_NAME = "EleutherAI/gpt-neo-125M"

TARGET_TOKENS = 39_000_000_000
CHUNK_BYTE_LIMIT = 250 * 1024 * 1024 
BATCH_SIZE = 1500
MAX_WORKERS = max(1, (os.cpu_count() or 4) - 1)

STATE_FILE = "titan_state.json"
DB_FILE = "exact_hashes.db"
LSH_FILE = "lsh_index.pkl"
BLOOM_FILE = "bloom_filter.pkl"

STATE = {
    "status": "Booting Singularity v12.0",
    "row_offset": 0,
    "total_tokens": 0,
    "stories_saved": 0,
    "current_chunk_id": 1,
    "dropped_quality": 0,
    "dropped_dupes": 0,
    "lsh_resets": 0,
    "bytes_written_manual": 0
}

STOPWORDS = {"the", "be", "to", "of", "and", "a", "in", "that", "have", "i", "it", "for", "not", "on", "with", "he", "as", "you", "do", "at", "this", "but", "his", "by", "from"}

tps_window = deque(maxlen=20)
state_lock = threading.Lock() # 🔱 Prevents race conditions across 16 cores

app = FastAPI()

class TitanEngine:
    def __init__(self):
        self.api = HfApi(token=HF_TOKEN)
        self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME, use_fast=True)
        self.cpu_pool = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
        self.io_pool = concurrent.futures.ThreadPoolExecutor(max_workers=2)
        
        self.upload_queue = asyncio.Queue(maxsize=5)
        self.db_write_queue = queue.Queue()
        
        self.last_state_save = time.time()
        self.last_heavy_save = time.time()
        
        self.init_db()
        self.load_state()
        
        self.cctx = zstd.ZstdCompressor(level=3)
        self.current_file = f"shard_{STATE['current_chunk_id']}.jsonl.zst"
        self._open_new_shard()

        threading.Thread(target=self.db_writer_worker, daemon=True).start()

    def _open_new_shard(self):
        self.shard_file_obj = open(self.current_file, "wb")
        self.zstd_stream = self.cctx.stream_writer(self.shard_file_obj)

    def _close_shard(self):
        self.zstd_stream.close()
        self.shard_file_obj.close()

    def init_db(self):
        self.read_conn = sqlite3.connect(DB_FILE, check_same_thread=False)
        self.read_cursor = self.read_conn.cursor()
        self.read_cursor.execute('PRAGMA journal_mode=WAL;')
        self.read_cursor.execute('PRAGMA synchronous=NORMAL;') # 🔱 Faster DB throughput
        self.read_cursor.execute('CREATE TABLE IF NOT EXISTS hashes (hash BLOB PRIMARY KEY)')
        self.read_conn.commit()
        self.db_read_lock = threading.Lock()

    def db_writer_worker(self):
        write_conn = sqlite3.connect(DB_FILE, check_same_thread=False)
        write_cursor = write_conn.cursor()
        write_cursor.execute('PRAGMA journal_mode=WAL;')
        write_cursor.execute('PRAGMA synchronous=NORMAL;')
        
        batch = []
        while True:
            try:
                item = self.db_write_queue.get(timeout=2.0)
                if item == "SHUTDOWN": break
                batch.append((item,))
            except queue.Empty:
                pass
            
            if len(batch) >= 500 or (batch and self.db_write_queue.empty()):
                write_cursor.executemany("INSERT OR IGNORE INTO hashes (hash) VALUES (?)", batch)
                write_conn.commit()
                batch.clear()

    def load_state(self):
        # 🚀 IMMORTAL SYNC 1: Try to load the JSON Master Ledger
        try:
            print("☁️ Checking Vault for Titan Ledger...")
            self.api.hf_hub_download(
                repo_id=PRIVATE_REPO, 
                repo_type="dataset", 
                filename=f"system_backups/{STATE_FILE}", 
                local_dir=".",
                force_download=True # 🛡️ Bypass broken local caches
            )
            if os.path.exists(STATE_FILE):
                with open(STATE_FILE, "r") as f: STATE.update(json.load(f))
            print("✅ Cloud Ledger loaded.")
        except Exception as e:
            print(f"⚠️ Ledger download failed: {e}. Relying on Vault Scan.")

        # 🚀 THE BULLETPROOF FALLBACK: Physically count the shards in the cloud
        print("🔍 Scanning Hugging Face Vault to protect existing data...")
        try:
            repo_files = self.api.list_repo_files(repo_id=PRIVATE_REPO, repo_type="dataset")
            shard_numbers = []
            for file in repo_files:
                # Look for 'silver_data/shard_54.jsonl.zst' etc.
                match = re.search(r'shard_(\d+)\.jsonl\.zst', file)
                if match:
                    shard_numbers.append(int(match.group(1)))
            
            if shard_numbers:
                max_shard = max(shard_numbers)
                print(f"📦 Vault Scan complete. Highest Titan shard found: {max_shard}")
                
                # 🛡️ Force the chunk ID so we NEVER overwrite good data
                if STATE["current_chunk_id"] <= max_shard:
                    STATE["current_chunk_id"] = max_shard + 1
                    print(f"🛑 OVERRIDE: Forcing next chunk ID to {STATE['current_chunk_id']}")
                
                # 🧠 Reconstruct estimated tokens if ledger was completely wiped
                if STATE["total_tokens"] < (max_shard * 50_000_000): 
                    estimated_tokens = max_shard * 55_000_000 # Rough estimate per Titan shard
                    STATE["total_tokens"] = estimated_tokens
                    print(f"📈 Recovered approx {estimated_tokens:,} tokens based on physical files.")
        except Exception as e:
            print(f"⚠️ Vault scan error: {e}")

        # Load Heavy Filters
        self.lsh = pickle.load(open(LSH_FILE, "rb")) if os.path.exists(LSH_FILE) else MinHashLSH(threshold=0.8, num_perm=128)
        self.bloom = pickle.load(open(BLOOM_FILE, "rb")) if os.path.exists(BLOOM_FILE) else ScalableBloomFilter(mode=ScalableBloomFilter.LARGE_SET_GROWTH)
    
    def save_state_atomic(self, heavy=False):
        with state_lock:
            state_copy = STATE.copy()
            
        with open(STATE_FILE + ".tmp", "w") as f: json.dump(state_copy, f)
        os.replace(STATE_FILE + ".tmp", STATE_FILE)
        
        # 🚀 IMMORTAL SYNC 2: Push the updated Master Ledger to the Cloud
        try:
            self.api.upload_file(
                path_or_fileobj=STATE_FILE, 
                path_in_repo=f"system_backups/{STATE_FILE}", 
                repo_id=PRIVATE_REPO, 
                repo_type="dataset"
            )
        except Exception as e:
            pass # If it fails, no worries. It will try again on the next loop.
            
        if heavy:
            with open(LSH_FILE + ".tmp", "wb") as f: pickle.dump(self.lsh, f)
            os.replace(LSH_FILE + ".tmp", LSH_FILE)
            with open(BLOOM_FILE + ".tmp", "wb") as f: pickle.dump(self.bloom, f)
            os.replace(BLOOM_FILE + ".tmp", BLOOM_FILE)

    def hard_filter(self, text):
        # 1. Drop tiny stubs or massive memory-hog strings
        if len(text) < 100 or len(text) > 500_000: 
            return None
        
        # 2. Drop blatant SEO spam
        if text.count("http") > 5: 
            return None
        
        # 3. Let the pure textbook data through
        return text

    def is_duplicate(self, text, tokens):
        h = xxhash.xxh64(text).digest() 
        
        if h in self.bloom:
            with self.db_read_lock:
                self.read_cursor.execute("SELECT 1 FROM hashes WHERE hash=?", (h,))
                if self.read_cursor.fetchone(): return True
        
        m = MinHash(num_perm=128)
        for gram in zip(tokens, tokens[1:], tokens[2:]):
            m.update(" ".join(gram).encode('utf8'))
            
        if self.lsh.query(m): return True
            
        self.bloom.add(h)
        self.db_write_queue.put(h) 
        self.lsh.insert(h.hex(), m)
        return False

    def process_batch(self, raw_batch):
        clean_batch, final_ids = [], []
        dropped_q, dropped_d = 0, 0
        
        for text in raw_batch:
            cleaned = self.hard_filter(text)
            if not cleaned: 
                dropped_q += 1; continue
                
            tokens = cleaned.split()
            if self.is_duplicate(cleaned, tokens):
                dropped_d += 1; continue
                
            clean_batch.append(cleaned)
            
        with state_lock:
            STATE["dropped_quality"] += dropped_q
            STATE["dropped_dupes"] += dropped_d
            
        if not clean_batch: return [], []
        
        ids = self.tokenizer(clean_batch, add_special_tokens=False, padding=False, truncation=False, return_attention_mask=False, return_token_type_ids=False)["input_ids"]
        return clean_batch, ids

    def write_zstd_batch(self, data_bytes):
        self.zstd_stream.write(data_bytes)
        # 🔱 Occasional flush prevents RAM bloat during massive runs
        if getattr(self, '_flush_counter', 0) > 50:
            self.zstd_stream.flush(zstd.FLUSH_BLOCK)
            self._flush_counter = 0
        else:
            self._flush_counter = getattr(self, '_flush_counter', 0) + 1

    @retry(stop=stop_after_attempt(5), wait=wait_random_exponential(multiplier=1, max=60))
    def robust_upload(self, filepath):
        print(f"🚀 UPLOADING SHARD: {filepath}")
        
        # 1. Upload the physical data shard
        self.api.upload_file(
            path_or_fileobj=filepath, 
            path_in_repo=f"silver_data/{filepath}", 
            repo_id=PRIVATE_REPO, 
            repo_type="dataset"
        )
        
        # 2. VERIFICATION SHIELD: Ask HF if the file is actually there
        if self.api.file_exists(repo_id=PRIVATE_REPO, filename=f"silver_data/{filepath}", repo_type="dataset"):
            print(f"✅ VERIFIED: {filepath} is safe in the vault.")
            
            # 3. Lock in the progress by forcing a cloud state update
            self.save_state_atomic(heavy=False)
            
            # 4. Only delete the local file after 100% cloud confirmation
            os.remove(filepath)
        else:
            raise ValueError(f"CRITICAL: Upload reported success but {filepath} is missing from the vault!")
    
    async def upload_worker(self):
        while True:
            filepath = await self.upload_queue.get()
            try: await asyncio.wait_for(asyncio.to_thread(self.robust_upload, filepath), timeout=600.0)
            except Exception as e: print(f"❌ Upload Fail: {e}")
            finally: self.upload_queue.task_done()

    async def run_ingestion(self):
        loop = asyncio.get_running_loop()
        dataset = load_dataset(DATASET_NAME, name="default", split="train", streaming=True).skip(STATE["row_offset"])
        it = iter(dataset)
        
        batch_raw = []
        start_time = time.time()
        tokens_in_window = 0
        
        with state_lock:
            STATE["status"] = "Engaged: Ingesting FineWeb-Edu 🚀" # 🔱 ADD THIS LINE!
            
        while STATE["total_tokens"] < TARGET_TOKENS:
          
            try:
                try: item = await asyncio.to_thread(next, it)
                except StopIteration: break               
                
                batch_raw.append(item.get("text", ""))
                
                with state_lock:
                    STATE["row_offset"] += 1
                
                if len(batch_raw) >= BATCH_SIZE:
                    texts, ids = await loop.run_in_executor(self.cpu_pool, self.process_batch, batch_raw)
                    batch_raw.clear()
                    
                    if texts:
                        batch_bytes = b"".join([(json.dumps({"text": t, "tok": len(d)}) + "\n").encode('utf-8') for t, d in zip(texts, ids)])
                        
                        await loop.run_in_executor(self.io_pool, self.write_zstd_batch, batch_bytes)
                        
                        tok_count = sum(len(d) for d in ids)
                        with state_lock:
                            STATE["total_tokens"] += tok_count
                            STATE["stories_saved"] += len(texts)
                            STATE["bytes_written_manual"] += len(batch_bytes)
                            current_bytes = STATE["bytes_written_manual"]
                            
                        tokens_in_window += tok_count
                    else:
                        with state_lock:
                            current_bytes = STATE["bytes_written_manual"]
                    
                    if time.time() - start_time >= 1.0:
                        tps_window.append(tokens_in_window / (time.time() - start_time))
                        start_time = time.time(); tokens_in_window = 0
                        
                    # 🚀 Throttled to 10 minutes (600 seconds) to avoid HF Commit Rate Limits
                    if time.time() - self.last_state_save >= 600.0:
                        await loop.run_in_executor(self.io_pool, self.save_state_atomic, False)
                        self.last_state_save = time.time()

                    # 🚀 Heavy filters now backup only once every 1 hour (3600 seconds)
                    if time.time() - self.last_heavy_save >= 3600.0:
                        await loop.run_in_executor(self.io_pool, self.save_state_atomic, True)
                        self.last_heavy_save = time.time()

                    if STATE["stories_saved"] % 5_000_000 == 0 and STATE["stories_saved"] > 0:
                        self.lsh = MinHashLSH(threshold=0.8, num_perm=128)
                        with state_lock:
                            STATE["lsh_resets"] += 1

                    if current_bytes >= CHUNK_BYTE_LIMIT:
                        await loop.run_in_executor(self.io_pool, self._close_shard)
                        await self.upload_queue.put(self.current_file)
                        
                        with state_lock:
                            STATE["current_chunk_id"] += 1
                            self.current_file = f"shard_{STATE['current_chunk_id']}.jsonl.zst"
                            STATE["bytes_written_manual"] = 0
                        
                        await loop.run_in_executor(self.io_pool, self._open_new_shard)
                        
            except asyncio.CancelledError:
                print("🛑 Ingestion paused cleanly for server reboot.")
                break
            except Exception as e:
                with state_lock:
                    STATE["status"] = f"CRITICAL CRASH: {e}"
                print(f"❌ CRASH TRACE: {e}")
                break
        
        # ==========================================
        # 🧹 BLOCK 5: THE TAIL-END SWEEP
        # ==========================================
        # If the loop finishes normally, catch the final partially-filled shard!
        if getattr(self, 'current_file', None) and STATE["bytes_written_manual"] > 0:
            print("🧹 Ingestion Complete! Sweeping the final partial shard into the vault...")
            with state_lock:
                STATE["status"] = "Finishing Up: Uploading final partial shard... 📦"
                
            # Close the final file and force it into the upload queue
            await loop.run_in_executor(self.io_pool, self._close_shard)
            await self.upload_queue.put(self.current_file)
            
            # Force one last heavy cloud backup
            await loop.run_in_executor(self.io_pool, self.save_state_atomic, True)
            
            with state_lock:
                STATE["status"] = "✅ MISSION ACCOMPLISHED: Target Tokens Reached!"

# --- 🛡️ V2 LIFESPAN: THE GRACEFUL SHUTDOWN ---
titan = None
@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    global titan
    titan = TitanEngine()
    
    # 1. Boot up the workers
    upload_tasks = [asyncio.create_task(titan.upload_worker()) for _ in range(3)]
    ingest_task = asyncio.create_task(titan.run_ingestion())
    
    yield # 🟢 The engine runs here while the server is alive
    
    # 🛑 HUGGING FACE SENT THE KILL SIGNAL (CPU RESTARTING)
    print("🚨 SIGTERM RECEIVED: Initiating Graceful Shutdown Sequence...")
    if titan:
        with state_lock:
            STATE["status"] = "Shutting Down: Draining Queues & Syncing Cloud..."
            
        # 2. Stop new data from processing
        ingest_task.cancel() 
        titan.db_write_queue.put("SHUTDOWN")
        
        # 🛡️ THE MISSING LOCK: Wait for the Tail-End Sweep to finish!
        try:
            await ingest_task 
        except asyncio.CancelledError:
            pass
        
        # 3. Force the upload queue to finish pushing any pending shards to the vault
        print(f"📦 Draining Upload Queue... ({titan.upload_queue.qsize()} files left)")
        await titan.upload_queue.join() 
        
        # 4. Perform the final atomic cloud sync of the Master Ledger & Heavy Filters
        print("☁️ Performing final Cloud Ledger Sync...")
        await asyncio.to_thread(titan.save_state_atomic, True)
        
        # 5. Safely kill the CPU threads
        titan.cpu_pool.shutdown(wait=False)
        titan.io_pool.shutdown(wait=False)
        print("✅ Shutdown Complete. Safe to reboot.")

app = FastAPI(lifespan=lifespan)


@app.get("/health")
async def health():
    avg_tps = sum(tps_window)/len(tps_window) if tps_window else 0
    with state_lock:
        st = STATE.copy()
    
    mb_written = st["bytes_written_manual"] / (1024 * 1024)
    return {
        "engine": "Indro-Nexus Titan v12.0 (The Singularity)",
        "status": st["status"], 
        "metrics": {
            "tokens": f"{st['total_tokens']:,} / {TARGET_TOKENS:,}", 
            "tps": f"{avg_tps:.0f} tk/s",
            "docs_saved": f"{st['stories_saved']:,}"
        },
        "storage": {
            "current_shard": st["current_chunk_id"],
            "shard_filling": f"{mb_written:.2f} MB / 250 MB",
            "upload_queue": titan.upload_queue.qsize() if titan else 0
        },
        "filters_dropped": {
            "quality": st["dropped_quality"],
            "duplicates": st["dropped_dupes"]
        }
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)