Add stream train autoloop script

scripts/stream_train_autoloop.py

@@ -0,0 +1,335 @@
+#!/usr/bin/env python3
+"""
+Stream egitim dongusu: chunk indir -> egit -> checkpoint/HF upload -> sonraki chunk.
+Takilirsa / cokurse otomatik resume.
+
+Ornek (VM):
+  export HF_TOKEN=...
+  export HF_REPO=HayrettinIscan/MeshAI-Base-Models
+  export HF_REPO_TYPE=model
+  python3 scripts/stream_train_autoloop.py --chunk-size 8 --epochs-per-chunk 1
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+CACHE = ROOT / "data" / "stream_cache"
+PROGRESS = ROOT / "data" / "stream_progress.json"
+CHECKPOINT_DIR = ROOT / "checkpoints"
+LOG_DIR = ROOT / "logs"
+HF_PREPROCESSED = os.getenv("HF_PREPROCESSED_REPO", "HayrettinIscan/MeshAI-Preprocessed-4K")
+HF_REPO = os.getenv("HF_REPO", "HayrettinIscan/MeshAI-Base-Models")
+HF_REPO_TYPE = os.getenv("HF_REPO_TYPE", "model")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+REQUIRED_FILES = ("geometry_latent.npz",)
+RENDER_CANDIDATES = tuple(
+    f"renders/{name}"
+    for i in range(4)
+    for name in (f"view_{i:02d}_tex.png", f"view_{i:02d}.png")
+)
+
+
+def _log(msg: str) -> None:
+    LOG_DIR.mkdir(parents=True, exist_ok=True)
+    line = f"[{datetime.now():%Y-%m-%d %H:%M:%S}] [stream] {msg}"
+    print(line, flush=True)
+    with open(LOG_DIR / "stream_train_autoloop.log", "a", encoding="utf-8") as f:
+        f.write(line + "\n")
+
+
+def _load_progress() -> dict:
+    if PROGRESS.exists():
+        try:
+            return json.loads(PROGRESS.read_text(encoding="utf-8"))
+        except Exception:
+            pass
+    return {"next_index": 0, "done_uids": [], "rounds": 0, "restarts": 0}
+
+
+def _save_progress(payload: dict) -> None:
+    PROGRESS.parent.mkdir(parents=True, exist_ok=True)
+    PROGRESS.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+
+def _load_manifest(token: str) -> list[dict]:
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(
+        repo_id=HF_PREPROCESSED,
+        filename="preprocessed/preprocessed_objects.json",
+        repo_type="dataset",
+        token=token,
+    )
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    rows = [r for r in payload.get("objects", []) if r.get("status", "ready") == "ready"]
+    if not rows:
+        raise RuntimeError("preprocessed_objects.json bos")
+    return rows
+
+
+def _uid_prefix(uid: str) -> str:
+    return f"preprocessed/{uid[:2]}/{uid}"
+
+
+def _copy_hub_file(cached: Path, target: Path) -> None:
+    target.parent.mkdir(parents=True, exist_ok=True)
+    if target.exists() and target.stat().st_size > 0:
+        return
+    shutil.copy2(cached, target)
+
+
+def _prefetch_uid(uid: str, token: str) -> bool:
+    """Tek obje dosyalarini stream_cache'e indir. Eksik/404 ise False."""
+    from huggingface_hub import hf_hub_download
+
+    dest = CACHE / uid[:2] / uid
+    dest.mkdir(parents=True, exist_ok=True)
+    try:
+        for name in REQUIRED_FILES:
+            cached = Path(
+                hf_hub_download(
+                    repo_id=HF_PREPROCESSED,
+                    filename=f"{_uid_prefix(uid)}/{name}",
+                    repo_type="dataset",
+                    token=token,
+                )
+            )
+            _copy_hub_file(cached, dest / name)
+        got_render = False
+        for rel in RENDER_CANDIDATES:
+            try:
+                cached = Path(
+                    hf_hub_download(
+                        repo_id=HF_PREPROCESSED,
+                        filename=f"{_uid_prefix(uid)}/{rel}",
+                        repo_type="dataset",
+                        token=token,
+                    )
+                )
+                _copy_hub_file(cached, dest / rel)
+                got_render = True
+            except Exception:
+                continue
+        if not got_render:
+            _log(f"UYARI: {uid} render yok, siyah pad ile devam")
+        return (dest / "geometry_latent.npz").exists()
+    except Exception as exc:
+        msg = str(exc).lower()
+        if "429" in msg or "rate limit" in msg:
+            _log(f"RATE LIMIT uid={uid}: {exc}")
+            raise
+        _log(f"ATLA uid={uid}: {exc}")
+        return False
+
+
+def _prune_cache(keep_uids: set[str]) -> None:
+    if not CACHE.exists():
+        return
+    for xx in CACHE.iterdir():
+        if not xx.is_dir():
+            continue
+        for uid_dir in xx.iterdir():
+            if uid_dir.is_dir() and uid_dir.name not in keep_uids:
+                shutil.rmtree(uid_dir, ignore_errors=True)
+
+
+def _upload_checkpoint() -> None:
+    latest = CHECKPOINT_DIR / "latest_model.pt"
+    if not latest.exists() or not HF_TOKEN:
+        return
+    try:
+        from huggingface_hub import HfApi
+
+        api = HfApi(token=HF_TOKEN)
+        api.upload_file(
+            path_or_fileobj=str(latest),
+            path_in_repo="checkpoints/latest_model.pt",
+            repo_id=HF_REPO,
+            repo_type=HF_REPO_TYPE,
+            token=HF_TOKEN,
+            commit_message="stream autoloop checkpoint",
+        )
+        _log(f"HF upload OK -> {HF_REPO}/checkpoints/latest_model.pt ({latest.stat().st_size // 1024} KB)")
+    except Exception as exc:
+        _log(f"HF upload basarisiz (devam): {exc}")
+
+
+def _run_train_chunk(epochs: int, stall_sec: int, checkpoint_every: int) -> int:
+    """Train subprocess; stall_sec boyunca log yoksa oldur. returncode dondur."""
+    cmd = [
+        sys.executable,
+        "-u",
+        str(ROOT / "train_pipeline.py"),
+        "--mode",
+        "real",
+        "--epochs",
+        str(epochs),
+        "--data-root",
+        str(CACHE),
+        "--validation-every",
+        "999999",
+        "--checkpoint-every",
+        str(checkpoint_every),
+    ]
+    latest = CHECKPOINT_DIR / "latest_model.pt"
+    if latest.exists():
+        cmd.extend(["--resume_from", str(latest)])
+
+    env = os.environ.copy()
+    env["PYTHONUNBUFFERED"] = "1"
+    _log(f"TRAIN: {' '.join(cmd)}")
+
+    proc = subprocess.Popen(
+        cmd,
+        cwd=str(ROOT),
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
+    assert proc.stdout is not None
+    last_out = time.time()
+    while True:
+        line = proc.stdout.readline()
+        if line:
+            print(line, end="", flush=True)
+            last_out = time.time()
+            if "checkpoint_saved" in line.lower():
+                _upload_checkpoint()
+            continue
+        if proc.poll() is not None:
+            break
+        if time.time() - last_out > stall_sec:
+            _log(f"STALL {stall_sec}s — process olduruluyor (otomatik restart)")
+            proc.kill()
+            try:
+                proc.wait(timeout=30)
+            except Exception:
+                pass
+            return 124
+        time.sleep(0.5)
+    return int(proc.returncode or 0)
+
+
+def run_once(args: argparse.Namespace) -> bool:
+    """Bir tam tur (manifest sonuna kadar). True=bitti, False=hata/yeniden dene."""
+    if not HF_TOKEN:
+        _log("HATA: HF_TOKEN yok")
+        return True
+
+    progress = _load_progress()
+    rows = _load_manifest(HF_TOKEN)
+    n = len(rows)
+    idx = int(progress.get("next_index", 0))
+    _log(f"Manifest={n} | next_index={idx} | chunk={args.chunk_size}")
+
+    while idx < n:
+        chunk = rows[idx : idx + args.chunk_size]
+        uids = [str(r.get("uid") or r.get("object_id")) for r in chunk]
+        _log(f"CHUNK [{idx}:{idx + len(uids)}] -> {uids[:3]}{'...' if len(uids) > 3 else ''}")
+
+        ok_uids: list[str] = []
+        for uid in uids:
+            try:
+                if _prefetch_uid(uid, HF_TOKEN):
+                    ok_uids.append(uid)
+                else:
+                    progress.setdefault("skipped", []).append(uid)
+            except Exception as exc:
+                if "429" in str(exc) or "rate limit" in str(exc).lower():
+                    _log("Rate limit — 90s bekle, sonra restart")
+                    _save_progress(progress)
+                    time.sleep(90)
+                    return False
+                progress.setdefault("skipped", []).append(uid)
+
+        if not ok_uids:
+            idx += len(chunk)
+            progress["next_index"] = idx
+            _save_progress(progress)
+            continue
+
+        _prune_cache(set(ok_uids))
+        code = _run_train_chunk(
+            epochs=args.epochs_per_chunk,
+            stall_sec=args.stall_sec,
+            checkpoint_every=args.checkpoint_every,
+        )
+        if code == 124:
+            progress["restarts"] = int(progress.get("restarts", 0)) + 1
+            _save_progress(progress)
+            _log("Stall restart — ayni chunk tekrar denenecek")
+            time.sleep(15)
+            return False
+        if code != 0:
+            progress["restarts"] = int(progress.get("restarts", 0)) + 1
+            _save_progress(progress)
+            _log(f"Train exit={code} — 20s sonra restart")
+            time.sleep(20)
+            return False
+
+        done = list(progress.get("done_uids", []))
+        done.extend(ok_uids)
+        progress["done_uids"] = done[-500:]  # dosya sismesin
+        idx += len(chunk)
+        progress["next_index"] = idx
+        progress["rounds"] = int(progress.get("rounds", 0)) + 1
+        _save_progress(progress)
+        _upload_checkpoint()
+        _log(f"Chunk OK | ilerleme {idx}/{n}")
+
+    _log("Tum objeler islendi.")
+    return True
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="MeshAI stream train autoloop")
+    p.add_argument("--chunk-size", type=int, default=8, help="Her turda kac obje indir+egit")
+    p.add_argument("--epochs-per-chunk", type=int, default=1)
+    p.add_argument("--checkpoint-every", type=int, default=10)
+    p.add_argument("--stall-sec", type=int, default=600, help="Log yoksa oldur (sn)")
+    p.add_argument("--max-restarts", type=int, default=500)
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
+    CACHE.mkdir(parents=True, exist_ok=True)
+    _log(
+        f"START chunk={args.chunk_size} epochs/chunk={args.epochs_per_chunk} "
+        f"stall={args.stall_sec}s repo={HF_PREPROCESSED}"
+    )
+    restarts = 0
+    while restarts <= args.max_restarts:
+        try:
+            finished = run_once(args)
+            if finished:
+                _log("DONE")
+                return
+            restarts += 1
+            _log(f"Auto-restart #{restarts}")
+        except KeyboardInterrupt:
+            _log("Kullanici durdurdu")
+            raise
+        except Exception as exc:
+            restarts += 1
+            _log(f"Beklenmeyen hata: {exc} — restart #{restarts}")
+            time.sleep(30)
+    _log("max-restarts asildi, cikis")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()