Spaces:

Neon-AI
/

Crawl

Runtime error

File size: 11,006 Bytes

"""
Common Crawl URL Extractor
--------------------------
Streams every URL from the latest Common Crawl crawl into chunked .gz files.
Resumable: skips shards whose output file already exists.

HuggingFace Spaces keep-alive: a lightweight HTTP server runs on port 7860
serving a live status page — HF sees it as an active app and won't kill it.

Output: cc_urls/urls_chunk_000.gz, urls_chunk_001.gz, ...
        Each file contains one URL per line, gzip-compressed.

Estimated totals (latest single crawl):
  - ~300 CDX shards
  - ~800 MB per shard (compressed download)
  - ~3–5 billion URLs total
  - ~240 GB total download bandwidth
"""

import gzip
import json
import os
import time
import threading
import requests
from http.server import BaseHTTPRequestHandler, HTTPServer

# ── Config ────────────────────────────────────────────────────────────────────

OUTPUT_DIR   = os.environ.get("OUTPUT_DIR", "./cc_urls")
COLLINFO_URL = "https://index.commoncrawl.org/collinfo.json"
BASE_URL     = "https://data.commoncrawl.org/"
STREAM_CHUNK = 1 * 1024 * 1024     # 1 MB read chunks while streaming
MAX_RETRIES  = 5
RETRY_DELAY  = 10                   # seconds between retries
PORT         = 7860                 # HF Spaces default port

# ── Shared state (written by scraper, read by HTTP handler) ──────────────────

state = {
    "crawl_id":     "starting...",
    "total_shards": 0,
    "current":      0,
    "skipped":      0,
    "failed":       0,
    "total_urls":   0,
    "current_shard": "—",
    "current_urls":  0,
    "status":       "initialising",
    "start_time":   time.time(),
    "shard_start":  time.time(),
}
state_lock = threading.Lock()


def get_state():
    with state_lock:
        return dict(state)


def set_state(**kwargs):
    with state_lock:
        state.update(kwargs)


# ── Keep-alive HTTP server ────────────────────────────────────────────────────

STATUS_HTML = """<!doctype html>
<html>
<head>
  <meta charset="utf-8">
  <meta http-equiv="refresh" content="10">
  <title>CC Extractor</title>
  <style>
    body {{ font-family: monospace; background: #0d1117; color: #c9d1d9; padding: 2rem; }}
    h2   {{ color: #58a6ff; margin-bottom: 1rem; }}
    .row {{ display: flex; gap: 3rem; margin-bottom: .5rem; }}
    .lbl {{ color: #8b949e; width: 160px; }}
    .val {{ color: #e6edf3; font-weight: bold; }}
    .bar-wrap {{ background: #21262d; border-radius: 4px; height: 16px; width: 400px; margin-top: 1rem; }}
    .bar      {{ background: #1f6feb; border-radius: 4px; height: 16px; transition: width .5s; }}
    .ok   {{ color: #3fb950; }}
    .fail {{ color: #f85149; }}
    .note {{ color: #8b949e; font-size: .85rem; margin-top: 2rem; }}
  </style>
</head>
<body>
  <h2>&#127758; Common Crawl URL Extractor</h2>
  {rows}
  <div class="bar-wrap"><div class="bar" style="width:{pct}%"></div></div>
  <p class="note">Page auto-refreshes every 10 s &nbsp;|&nbsp; Started {elapsed}</p>
</body>
</html>"""


def make_status_html():
    s = get_state()
    elapsed_s = int(time.time() - s["start_time"])
    h, rem = divmod(elapsed_s, 3600)
    m, sec = divmod(rem, 60)
    elapsed_str = f"{h}h {m}m {sec}s"

    pct = 0
    if s["total_shards"] > 0:
        pct = round(s["current"] / s["total_shards"] * 100, 1)

    shard_elapsed = int(time.time() - s["shard_start"])

    def row(label, value, cls="val"):
        return f'<div class="row"><span class="lbl">{label}</span><span class="{cls}">{value}</span></div>'

    rows = "\n".join([
        row("Crawl",          s["crawl_id"]),
        row("Status",         s["status"]),
        row("Progress",       f'{s["current"]} / {s["total_shards"]} shards  ({pct}%)'),
        row("Current shard",  s["current_shard"]),
        row("Shard URLs",     f'{s["current_urls"]:,}  ({shard_elapsed}s so far)'),
        row("Total URLs",     f'{s["total_urls"]:,}'),
        row("Skipped",        f'{s["skipped"]:,}', "ok"),
        row("Failed",         f'{s["failed"]:,}',  "fail" if s["failed"] else "ok"),
    ])

    return STATUS_HTML.format(rows=rows, pct=pct, elapsed=elapsed_str)


class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        body = make_status_html().encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *_):
        pass  # silence access logs


def start_server():
    server = HTTPServer(("0.0.0.0", PORT), Handler)
    t = threading.Thread(target=server.serve_forever, daemon=True)
    t.start()
    print(f"Status page running on http://0.0.0.0:{PORT}")


# ── Scraper helpers ───────────────────────────────────────────────────────────

def get_latest_crawl_paths():
    set_state(status="fetching crawl list")
    print("Fetching crawl list from collinfo.json ...")
    info = requests.get(COLLINFO_URL, timeout=30).json()
    latest_id = info[0]["id"]
    paths_url = f"{BASE_URL}crawl-data/{latest_id}/cc-index.paths.gz"
    print(f"Latest crawl: {latest_id}")

    r = requests.get(paths_url, timeout=60)
    r.raise_for_status()

    raw = gzip.decompress(r.content).decode("utf-8")
    paths = [
        line.strip()
        for line in raw.splitlines()
        if line.strip().endswith(".gz") and "/indexes/cdx-" in line
    ]
    print(f"Found {len(paths)} CDX shards.")
    set_state(crawl_id=latest_id, total_shards=len(paths))
    return latest_id, paths


def shard_output_path(shard_index):
    return os.path.join(OUTPUT_DIR, f"urls_chunk_{shard_index:03d}.gz")


def extract_urls_from_shard(shard_url, out_path):
    count = 0
    leftover = b""

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            r = requests.get(shard_url, stream=True, timeout=120)
            r.raise_for_status()

            with gzip.open(out_path + ".tmp", "wb") as out_gz:
                decompressor = gzip.GzipFile(fileobj=r.raw)
                while True:
                    chunk = decompressor.read(STREAM_CHUNK)
                    if not chunk:
                        break
                    block = leftover + chunk
                    lines = block.split(b"\n")
                    leftover = lines[-1]
                    for line in lines[:-1]:
                        line = line.strip()
                        if not line:
                            continue
                        try:
                            parts = line.split(b" ", 2)
                            if len(parts) < 3:
                                continue
                            obj = json.loads(parts[2])
                            url = obj.get("url")
                            if url:
                                out_gz.write((url + "\n").encode("utf-8"))
                                count += 1
                                # Update live counter every 50k URLs
                                if count % 50_000 == 0:
                                    set_state(current_urls=count)
                        except (json.JSONDecodeError, UnicodeDecodeError):
                            continue

                if leftover.strip():
                    try:
                        parts = leftover.split(b" ", 2)
                        if len(parts) >= 3:
                            obj = json.loads(parts[2])
                            url = obj.get("url")
                            if url:
                                out_gz.write((url + "\n").encode("utf-8"))
                                count += 1
                    except Exception:
                        pass

            os.replace(out_path + ".tmp", out_path)
            return count

        except Exception as e:
            print(f"  Attempt {attempt}/{MAX_RETRIES} failed: {e}")
            if os.path.exists(out_path + ".tmp"):
                os.remove(out_path + ".tmp")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY * attempt)
            else:
                print(f"  Giving up on this shard after {MAX_RETRIES} attempts.")
                return 0


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    start_server()

    crawl_id, shard_paths = get_latest_crawl_paths()

    total_urls = 0
    skipped    = 0
    failed     = 0
    start_time = time.time()

    print(f"\nStarting extraction → {OUTPUT_DIR}/\n{'─'*60}")
    set_state(status="running", start_time=start_time)

    for i, shard_path in enumerate(shard_paths):
        out_path  = shard_output_path(i)
        shard_url = BASE_URL + shard_path
        shard_name = os.path.basename(shard_path)

        if os.path.exists(out_path):
            skipped += 1
            set_state(current=i + 1, skipped=skipped, current_shard=f"{shard_name} (skipped)")
            print(f"[{i+1:03d}/{len(shard_paths)}] SKIP  {shard_name}")
            continue

        set_state(
            current=i + 1,
            current_shard=shard_name,
            current_urls=0,
            shard_start=time.time(),
            status="downloading",
        )
        print(f"[{i+1:03d}/{len(shard_paths)}] Fetching {shard_name} ...")
        t0 = time.time()

        count = extract_urls_from_shard(shard_url, out_path)
        elapsed = time.time() - t0

        if count > 0:
            total_urls += count
            size_mb = os.path.getsize(out_path) / 1024**2
            set_state(total_urls=total_urls, current_urls=count)
            print(f"           ✓ {count:,} URLs  |  {size_mb:.1f} MB  |  {elapsed:.0f}s")
        else:
            failed += 1
            set_state(failed=failed)
            print(f"           ✗ Failed — shard skipped")

    set_state(status="done", current_shard="—", current_urls=0)

    total_elapsed = time.time() - start_time
    print(f"\n{'─'*60}")
    print(f"Crawl:      {crawl_id}")
    print(f"Shards:     {len(shard_paths)} total  |  {skipped} skipped  |  {failed} failed")
    print(f"URLs:       {total_urls:,}")
    print(f"Output dir: {OUTPUT_DIR}/")
    print(f"Time:       {total_elapsed/3600:.2f} hours")
    print(f"{'─'*60}")

    # Keep server alive after completion so HF doesn't kill the space
    print("Extraction complete. Server staying up — safe to stop manually.")
    while True:
        time.sleep(60)


if __name__ == "__main__":
    main()