Update app.py

app.py

@@ -3,23 +3,21 @@ import json
 import time
 import socket
 import threading
-import io
+import re
 import requests
 from pathlib import Path
-from huggingface_hub import HfApi, list_repo_tree
+from huggingface_hub import HfApi
 
 # ── Config ───────────────────────────────────────────────────────────────────
-HF_TOKEN      = os.environ.get("HF_TOKEN")
-DATASET_REPO  = "HuggingFaceFW/fineweb-edu"
-RAW_DIR       = "/data/raw"
-STATE_FILE    = "/data/state.json"
-WORKER_TIMEOUT = 600  # 10 min — reclaim stale claimed shards
-
-# CC-MAIN-2025 prefix filter
-CC_PREFIX = "data/CC-MAIN-2025"
+HF_TOKEN       = os.environ.get("HF_TOKEN")
+DATASET_REPO   = "HuggingFaceFW/fineweb-edu"
+RAW_DIR        = "/data/raw"
+STATE_FILE     = "/data/state.json"
+WORKER_TIMEOUT = 600
+MAX_BUFFERED   = 4
+CC_PREFIX      = "data/CC-MAIN-2025"
 
 os.makedirs(RAW_DIR, exist_ok=True)
-
 api = HfApi(token=HF_TOKEN)
 
 # ── Keep-alive ────────────────────────────────────────────────────────────────
@@ -34,18 +32,27 @@ def serve():
         conn.send(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK")
         conn.close()
 
+# ── Friendly name ─────────────────────────────────────────────────────────────
+# data/CC-MAIN-2025-05/000_00042.parquet → cc2025-05_000042.parquet
+def friendly_name(hf_path):
+    m = re.search(r"CC-MAIN-(\d{4}-\d+)/\d+_(\d+)\.parquet", hf_path)
+    if m:
+        return f"cc{m.group(1)}_{int(m.group(2)):06d}.parquet"
+    return hf_path.replace("/", "__")
+
 # ── State ─────────────────────────────────────────────────────────────────────
 def load_state():
     if os.path.exists(STATE_FILE):
         with open(STATE_FILE) as f:
             state = json.load(f)
-        total   = len(state["shards"])
-        done    = sum(1 for s in state["shards"].values() if s["status"] == "done")
-        claimed = sum(1 for s in state["shards"].values() if s["status"] == "claimed")
-        pending = sum(1 for s in state["shards"].values() if s["status"] == "pending")
-        print(f"Resuming — {done} done / {claimed} claimed / {pending} pending / {total} total")
+        shards  = state["shards"]
+        queue   = state.get("queue", [])
+        done    = sum(1 for v in shards.values() if v["status"] == "done")
+        claimed = sum(1 for v in shards.values() if v["status"] == "claimed")
+        pending = sum(1 for v in shards.values() if v["status"] == "pending")
+        print(f"Resuming — {done} done / {claimed} claimed / {pending} buffered / {len(queue)} queued")
     else:
-        state = {"shards": {}}
+        state = {"shards": {}, "queue": []}
         print("Starting fresh")
     return state
 
@@ -55,88 +62,70 @@ def save_state(state):
         json.dump(state, f, indent=2)
     os.replace(tmp, STATE_FILE)
 
-# ── Discover all CC-MAIN-2025 parquet files ───────────────────────────────────
-def discover_shards(state):
+# ── Discover — builds queue only, does NOT touch state["shards"] ──────────────
+def discover_queue(state):
     print("Discovering shards from HF...")
     files = api.list_repo_files(DATASET_REPO, repo_type="dataset")
+    known = {v["hf_path"] for v in state["shards"].values()} | set(state.get("queue", []))
     new_count = 0
     for f in files:
-        if f.startswith(CC_PREFIX) and f.endswith(".parquet"):
-            if f not in state["shards"]:
-                state["shards"][f] = {
-                    "status":     "pending",
-                    "worker":     None,
-                    "claimed_at": None,
-                }
-                new_count += 1
-    print(f"✓ {new_count} new shards discovered | {len(state['shards'])} total")
+        if f.startswith(CC_PREFIX) and f.endswith(".parquet") and f not in known:
+            state["queue"].append(f)
+            new_count += 1
+    print(f"✓ {new_count} queued | {len(state['queue'])} in queue | {len(state['shards'])} in state")
     save_state(state)
 
 # ── Reclaim timed-out shards ──────────────────────────────────────────────────
 def reclaim_stale(state):
     now = time.time()
     reclaimed = 0
-    for shard, info in state["shards"].items():
+    for name, info in state["shards"].items():
         if info["status"] == "claimed" and info["claimed_at"]:
             if now - info["claimed_at"] > WORKER_TIMEOUT:
-                print(f"  ⚠ Reclaiming stale shard: {shard} (worker: {info['worker']})")
+                print(f"  ⚠ Reclaiming: {name} (worker: {info['worker']})")
                 info["status"]     = "pending"
                 info["worker"]     = None
                 info["claimed_at"] = None
                 reclaimed += 1
     if reclaimed:
         save_state(state)
-    return reclaimed
 
-# ── Download pending shards to /data/raw ─────────────────────────────────────
+# ── Download loop — pops queue, downloads, THEN adds to state as pending ──────
 def download_loop(state):
     base_url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/"
 
     while True:
-        # Reclaim stale first
-        reclaim_stale(state)
-
-        # Reload state to pick up worker updates
-        if os.path.exists(STATE_FILE):
+        try:
             with open(STATE_FILE) as f:
-                state["shards"] = json.load(f)["shards"]
+                fresh = json.load(f)
+            state["shards"] = fresh["shards"]
+            state["queue"]  = fresh.get("queue", [])
+        except Exception:
+            pass
 
-        # Count how many raw files already sitting in /data/raw (not yet claimed)
-        raw_files = list(Path(RAW_DIR).glob("*.parquet"))
-        pending_raw = len(raw_files)
+        reclaim_stale(state)
 
-        # Keep at most 4 shards pre-downloaded to avoid filling disk
-        if pending_raw >= 4:
-            print(f"  Buffer full ({pending_raw} shards waiting) — sleeping...")
-            time.sleep(60)
+        buffered = sum(1 for v in state["shards"].values() if v["status"] == "pending")
+        if buffered >= MAX_BUFFERED:
+            time.sleep(30)
             continue
 
-        # Find next pending shard to download
-        to_download = None
-        for shard, info in state["shards"].items():
-            if info["status"] == "pending":
-                raw_name = shard.replace("/", "__")
-                raw_path = Path(RAW_DIR) / raw_name
-                if not raw_path.exists():
-                    to_download = shard
-                    break
-
-        if not to_download:
-            done  = sum(1 for s in state["shards"].values() if s["status"] == "done")
+        if not state["queue"]:
+            done  = sum(1 for v in state["shards"].values() if v["status"] == "done")
             total = len(state["shards"])
-            if done == total:
+            if done == total and total > 0:
                 print("✓ All shards complete!")
                 break
-            print("  Nothing to download right now — sleeping...")
+            print("  Queue empty — sleeping...")
             time.sleep(60)
             continue
 
-        # Download it
-        url      = base_url + to_download
-        raw_name = to_download.replace("/", "__")
-        raw_path = Path(RAW_DIR) / raw_name
+        hf_path  = state["queue"][0]
+        name     = friendly_name(hf_path)
+        raw_path = Path(RAW_DIR) / name
+        url      = base_url + hf_path
 
-        print(f"  Downloading: {to_download}")
+        print(f"  Downloading: {hf_path} → {name}")
         try:
             resp = requests.get(
                 url,
@@ -148,37 +137,47 @@ def download_loop(state):
             with open(raw_path, "wb") as f:
                 for chunk in resp.iter_content(chunk_size=8 * 1024 * 1024):
                     f.write(chunk)
-            print(f"  ✓ Downloaded: {raw_name}")
         except Exception as e:
-            print(f"  ✗ Download failed: {e}")
+            print(f"  ✗ Failed: {e} — retrying in 30s")
             time.sleep(30)
             continue
 
+        # Only now: pop from queue, add to state as pending
+        state["queue"].pop(0)
+        state["shards"][name] = {
+            "status":     "pending",
+            "hf_path":    hf_path,
+            "worker":     None,
+            "claimed_at": None,
+        }
+        save_state(state)
+        print(f"  ✓ Ready: {name}")
         time.sleep(5)
 
-# ── Monitor loop — prints progress ───────────────────────────────────────────
-def monitor_loop(state):
+# ── Monitor ───────────────────────────────────────────────────────────────────
+def monitor_loop():
     while True:
         time.sleep(120)
-        if os.path.exists(STATE_FILE):
+        try:
             with open(STATE_FILE) as f:
-                s = json.load(f)["shards"]
-            done    = sum(1 for v in s.values() if v["status"] == "done")
-            claimed = sum(1 for v in s.values() if v["status"] == "claimed")
-            pending = sum(1 for v in s.values() if v["status"] == "pending")
-            total   = len(s)
+                s = json.load(f)
+            shards  = s["shards"]
+            queue   = s.get("queue", [])
+            done    = sum(1 for v in shards.values() if v["status"] == "done")
+            claimed = sum(1 for v in shards.values() if v["status"] == "claimed")
+            pending = sum(1 for v in shards.values() if v["status"] == "pending")
+            total   = len(shards) + len(queue)
             pct     = (done / total * 100) if total else 0
-            print(f"[MONITOR] {done}/{total} done ({pct:.1f}%) | {claimed} active | {pending} pending")
+            print(f"[MONITOR] {done}/{total} ({pct:.1f}%) | {claimed} active | {pending} buffered | {len(queue)} queued")
+        except Exception:
+            pass
 
 # ── Entry point ───────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     threading.Thread(target=serve, daemon=True).start()
-
     state = load_state()
-    discover_shards(state)
-
-    threading.Thread(target=monitor_loop, args=(state,), daemon=True).start()
+    discover_queue(state)
+    threading.Thread(target=monitor_loop, daemon=True).start()
     threading.Thread(target=download_loop, args=(state,), daemon=True).start()
-
     while True:
         time.sleep(60)