AdhyanshVerma's picture
Upload 16 files
207535f verified
Raw
History Blame Contribute Delete
26.1 kB
#!/usr/bin/env python3
"""
YT Pipeline Orchestrator v3 β€” Host System (Part 2)
===================================================
Kaggle limits: max 10 concurrent sessions, 12hr max runtime.
Auto-rotates workers, syncs results, retries failed tasks.
Usage:
python orchestrator.py run --videos-per-worker 3 # auto-pilot loop
python orchestrator.py launch --max-workers 5 # one-shot launch
python orchestrator.py status
python orchestrator.py sync
python orchestrator.py retry
python orchestrator.py import
"""
import argparse, json, logging, os, shutil, sqlite3, subprocess, sys
import time, uuid
from datetime import datetime, timezone, timedelta
from pathlib import Path
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
datefmt="%H:%M:%S",
)
log = logging.getLogger("orchestrator")
# ── paths & constants ────────────────────────────────────────────────
BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = Path(os.environ.get("DATA_DIR", "/data"))
if not DATA_DIR.exists():
try:
DATA_DIR.mkdir(parents=True, exist_ok=True)
except PermissionError:
DATA_DIR = BASE_DIR / "db"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DB_ACCOUNTS = DATA_DIR / "accounts.db"
DB_CHANNELS = DATA_DIR / "channel_links.db"
DB_TRACKER = DATA_DIR / "tracker.db"
WORKERS_DIR = BASE_DIR / "kaggle_workers"
WORKER_SCRIPT = Path(__file__).resolve().parent / "worker.py"
HF_REPO = "AdhyanshVerma/YT"
# kaggle limits
MAX_KAGGLE_SESSIONS = 5
KAGGLE_TIMEOUT_HOURS = 12
POLL_INTERVAL_SEC = 120 # check every 2 min in run loop
SAFETY_MARGIN_MIN = 30 # requeue if within 30min of 12hr limit
# video states
S_PENDING = "pending"
S_ASSIGNED = "assigned"
S_DONE = "done"
S_FAILED = "failed"
S_TIMEOUT = "timeout"
# worker states
WS_CREATED = "created"
WS_PUSHED = "pushed"
WS_RUNNING = "running"
WS_COMPLETE = "complete"
WS_ERROR = "error"
WS_TIMEOUT = "timeout"
WS_CANCELLED = "cancelled"
# ── tracker database ─────────────────────────────────────────────────
class Tracker:
def __init__(self, db_path=DB_TRACKER):
self.db_path = str(db_path)
self._init_db()
def _conn(self):
c = sqlite3.connect(self.db_path)
c.row_factory = sqlite3.Row
c.execute("PRAGMA journal_mode=WAL")
return c
def _init_db(self):
with self._conn() as c:
c.execute("""CREATE TABLE IF NOT EXISTS videos (
video_id TEXT PRIMARY KEY, url TEXT NOT NULL,
title TEXT DEFAULT '', channel TEXT DEFAULT '',
duration REAL DEFAULT 0, status TEXT DEFAULT 'pending',
worker_id TEXT DEFAULT '', assigned_at TEXT DEFAULT '',
completed_at TEXT DEFAULT '', error_msg TEXT DEFAULT '',
retry_count INTEGER DEFAULT 0
)""")
c.execute("""CREATE TABLE IF NOT EXISTS workers (
worker_id TEXT PRIMARY KEY, kaggle_title TEXT DEFAULT '',
video_count INTEGER DEFAULT 0, status TEXT DEFAULT 'created',
created_at TEXT DEFAULT '', pushed_at TEXT DEFAULT '',
finished_at TEXT DEFAULT '', kaggle_status TEXT DEFAULT ''
)""")
c.execute("CREATE INDEX IF NOT EXISTS idx_vid_status ON videos(status)")
c.execute("CREATE INDEX IF NOT EXISTS idx_w_status ON workers(status)")
# migrate old tables missing kaggle_status
try:
c.execute("ALTER TABLE workers ADD COLUMN kaggle_status TEXT DEFAULT ''")
except Exception:
pass # column already exists
def import_from_channels_db(self):
if not DB_CHANNELS.exists():
return 0
src = sqlite3.connect(str(DB_CHANNELS))
rows = src.execute("SELECT video_id,url,title,channel_url,duration FROM videos").fetchall()
src.close()
count = 0
with self._conn() as c:
for r in rows:
try:
c.execute("INSERT OR IGNORE INTO videos (video_id,url,title,channel,duration,status) VALUES (?,?,?,?,?,?)",
(r[0], r[1], r[2] or "", r[3] or "", r[4] or 0, S_PENDING))
count += 1
except: pass
log.info(f"Imported {count} videos")
return count
def get_pending(self, limit):
with self._conn() as c:
return c.execute("SELECT video_id,url,title FROM videos WHERE status=? ORDER BY duration ASC LIMIT ?",
(S_PENDING, limit)).fetchall()
def assign_batch(self, video_ids, worker_id):
now = datetime.now(timezone.utc).isoformat()
with self._conn() as c:
for vid in video_ids:
c.execute("UPDATE videos SET status=?,worker_id=?,assigned_at=? WHERE video_id=?",
(S_ASSIGNED, worker_id, now, vid))
def requeue_worker_videos(self, worker_id, new_status=S_PENDING):
"""Put assigned videos from a dead/timed-out worker back to pending."""
with self._conn() as c:
r = c.execute("UPDATE videos SET status=?,worker_id='' WHERE worker_id=? AND status=?",
(new_status, worker_id, S_ASSIGNED))
return r.rowcount
def mark_worker(self, worker_id, status, kaggle_status=""):
now = datetime.now(timezone.utc).isoformat()
with self._conn() as c:
if status in (WS_COMPLETE, WS_ERROR, WS_TIMEOUT, WS_CANCELLED):
c.execute("UPDATE workers SET status=?,finished_at=?,kaggle_status=? WHERE worker_id=?",
(status, now, kaggle_status, worker_id))
else:
c.execute("UPDATE workers SET status=?,kaggle_status=? WHERE worker_id=?",
(status, kaggle_status, worker_id))
def update_video_statuses(self, status_dict):
with self._conn() as c:
for vid, st in status_dict.items():
if st in (S_DONE, S_FAILED):
c.execute("UPDATE videos SET status=? WHERE video_id=? AND status!=?",
(st, vid, st))
def register_worker(self, worker_id, kaggle_title, video_count):
with self._conn() as c:
c.execute("INSERT OR REPLACE INTO workers VALUES (?,?,?,?,?,?,?,?)",
(worker_id, kaggle_title, video_count, WS_CREATED,
datetime.now(timezone.utc).isoformat(), "", "", ""))
def mark_worker_pushed(self, worker_id):
with self._conn() as c:
c.execute("UPDATE workers SET status=?,pushed_at=? WHERE worker_id=?",
(WS_PUSHED, datetime.now(timezone.utc).isoformat(), worker_id))
def active_workers(self):
"""Workers that are pushed/running (consuming Kaggle slots)."""
with self._conn() as c:
return c.execute("SELECT * FROM workers WHERE status IN (?,?)",
(WS_PUSHED, WS_RUNNING)).fetchall()
def pushed_workers(self):
with self._conn() as c:
return c.execute("SELECT * FROM workers WHERE status IN (?,?)",
(WS_PUSHED, WS_RUNNING)).fetchall()
def timed_out_workers(self):
"""Workers pushed more than 12hrs ago still active."""
cutoff = (datetime.now(timezone.utc) - timedelta(hours=KAGGLE_TIMEOUT_HOURS,
minutes=-SAFETY_MARGIN_MIN)).isoformat()
with self._conn() as c:
return c.execute("SELECT * FROM workers WHERE status IN (?,?) AND pushed_at<? AND pushed_at!=''",
(WS_PUSHED, WS_RUNNING, cutoff)).fetchall()
def retry_failed(self, max_retries=3):
with self._conn() as c:
r = c.execute("UPDATE videos SET status=?, retry_count=retry_count+1 WHERE status IN (?,?) AND retry_count<?",
(S_PENDING, S_FAILED, S_TIMEOUT, max_retries))
return r.rowcount
def stats(self):
with self._conn() as c:
rows = c.execute("SELECT status,COUNT(*) FROM videos GROUP BY status").fetchall()
return {r[0]: r[1] for r in rows}
def worker_stats(self):
with self._conn() as c:
return c.execute("SELECT * FROM workers ORDER BY created_at DESC LIMIT 20").fetchall()
def available_slots(self):
active = len(self.active_workers())
return max(0, MAX_KAGGLE_SESSIONS - active)
# ── helpers ──────────────────────────────────────────────────────────
def load_api_keys():
conn = sqlite3.connect(str(DB_ACCOUNTS))
rows = conn.execute("SELECT api_key FROM accounts WHERE api_key IS NOT NULL AND api_key!=''").fetchall()
conn.close()
return [r[0] for r in rows]
def load_hf_token():
token = os.environ.get("HF_TOKEN", "")
if not token:
p = Path.home() / ".cache" / "huggingface" / "token"
if p.exists(): token = p.read_text().strip()
return token
def get_kaggle_username():
kf = Path.home() / ".kaggle" / "kaggle.json"
if kf.exists(): return json.loads(kf.read_text()).get("username", "adhyanshverma")
return "adhyanshverma"
def check_kaggle_kernel_status(kaggle_title, username):
"""Query Kaggle API for kernel status. Returns: queued/running/complete/error/cancelled or None."""
try:
r = subprocess.run(["kaggle","kernels","status",f"{username}/{kaggle_title}"],
capture_output=True, text=True, timeout=30)
out = r.stdout.strip().lower()
for s in ["complete", "error", "cancelled", "running", "queued"]:
if s in out: return s
return out if out else None
except Exception:
return None
# ── notebook generation ──────────────────────────────────────────────
def generate_notebook(worker_dir, config, worker_code):
config_json = json.dumps(config, indent=2)
cookies_txt = Path("cookies.txt").read_text() if Path("cookies.txt").exists() else ""
cookies_json_str = Path("cookies.json").read_text() if Path("cookies.json").exists() else ""
cells = [
{"cell_type":"code","execution_count":None,"metadata":{},"outputs":[],
"source":["!pip install -q yt-dlp huggingface_hub openai faster-whisper pyarrow\n",
"!apt-get install -y -qq ffmpeg > /dev/null 2>&1\n",
"print('Dependencies installed βœ…')\n"]},
{"cell_type":"code","execution_count":None,"metadata":{},"outputs":[],
"source":["from pathlib import Path\n",
f"Path('cookies.txt').write_text({repr(cookies_txt)})\n",
f"Path('cookies.json').write_text({repr(cookies_json_str)})\n",
"print('Cookies written βœ…')\n"]},
{"cell_type":"code","execution_count":None,"metadata":{},"outputs":[],
"source": worker_code.split("\n")},
{"cell_type":"code","execution_count":None,"metadata":{},"outputs":[],
"source":["import json, os\n",
f"config = json.loads('''{config_json}''')\n",
"os.environ['HF_TOKEN'] = config['hf_token']\n",
"status = run_worker(config)\n",
"print(f'Worker finished: {json.dumps(status, indent=2)}')\n"]},
]
for cell in cells:
cell["source"] = [l if l.endswith("\n") else l+"\n" for l in cell["source"]]
nb = {"cells":cells,"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},
"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":4}
(worker_dir/"notebook.ipynb").write_text(json.dumps(nb, indent=2))
def generate_kernel_metadata(worker_dir, kaggle_title, username):
meta = {"id":f"{username}/{kaggle_title}","title":kaggle_title,"code_file":"notebook.ipynb",
"language":"python","kernel_type":"notebook","is_private":True,
"enable_gpu":False,"enable_internet":True,
"dataset_sources":[],"competition_sources":[],"kernel_sources":[]}
(worker_dir/"kernel-metadata.json").write_text(json.dumps(meta, indent=2))
def load_vision_models():
VISION_MODELS_FILE = BASE_DIR / "models" / "vision_models.txt"
if VISION_MODELS_FILE.exists():
with open(VISION_MODELS_FILE) as f:
models = [l.strip() for l in f if l.strip() and not l.strip().startswith("#")]
if models:
return models
return [
"MiniMaxAI/MiniMax-M3",
"XiaomiMiMo/MiMo-V2.5",
"Qwen/Qwen3-VL-32B-Instruct",
"Qwen/Qwen3.6-27B",
"Qwen/Qwen3-VL-235B-A22B-Thinking"
]
# ── launch N workers (respecting slot limit) ─────────────────────────
def launch_workers(tracker, num_workers, videos_per_worker, push=False):
slots = tracker.available_slots()
if slots <= 0:
log.info(f"No Kaggle slots available (all {MAX_KAGGLE_SESSIONS} in use)")
return 0
num_workers = min(num_workers, slots)
total_fetch = num_workers * videos_per_worker
pending = tracker.get_pending(total_fetch)
if not pending:
log.info("No pending videos")
return 0
api_keys = load_api_keys()
hf_token = load_hf_token()
if not api_keys or not hf_token:
log.error("Missing API keys or HF token")
return 0
username = get_kaggle_username()
worker_code = WORKER_SCRIPT.read_text()
chunks = [pending[i:i+videos_per_worker] for i in range(0, len(pending), videos_per_worker)]
WORKERS_DIR.mkdir(parents=True, exist_ok=True)
keys_per_w = max(1, len(api_keys) // max(1, len(chunks)))
launched = 0
vmodels = load_vision_models()
for widx, chunk in enumerate(chunks):
worker_id = f"w-{uuid.uuid4().hex[:8]}"
kaggle_title = f"yt-w-{uuid.uuid4().hex[:6]}"
start_k = (widx * keys_per_w) % len(api_keys)
wkeys = [api_keys[(start_k+k) % len(api_keys)] for k in range(keys_per_w)]
video_list = [{"video_id":v["video_id"],"url":v["url"]} for v in chunk]
video_ids = [v["video_id"] for v in chunk]
config = {"worker_id":worker_id,"videos":video_list,"api_keys":wkeys,
"hf_token":hf_token,"vision_models":vmodels,
"text_model":"Qwen/Qwen3-32B"}
wdir = WORKERS_DIR / worker_id
wdir.mkdir(parents=True, exist_ok=True)
generate_notebook(wdir, config, worker_code)
generate_kernel_metadata(wdir, kaggle_title, username)
tracker.assign_batch(video_ids, worker_id)
tracker.register_worker(worker_id, kaggle_title, len(chunk))
log.info(f" Worker: {kaggle_title} β€” {len(chunk)} videos, {len(wkeys)} keys")
if push:
r = subprocess.run(["kaggle","kernels","push","-p",str(wdir)],
capture_output=True, text=True)
if r.returncode == 0:
tracker.mark_worker_pushed(worker_id)
launched += 1
log.info(f" βœ… Pushed")
else:
log.error(f" ❌ Push failed: {r.stderr[:200]}")
tracker.requeue_worker_videos(worker_id)
tracker.mark_worker(worker_id, WS_ERROR, "push_failed")
else:
launched += 1
return launched
# ── sync: poll kaggle + HF status ───────────────────────────────────
def sync_workers(tracker):
"""Check Kaggle status for all active workers, handle timeouts."""
username = get_kaggle_username()
active = tracker.pushed_workers()
if not active:
return
log.info(f"Checking {len(active)} active workers...")
# 1. Check for 12hr timeouts
timed_out = tracker.timed_out_workers()
for w in timed_out:
wid = w["worker_id"]
log.warning(f" ⏰ Worker {wid} ({w['kaggle_title']}) exceeded 12hr limit")
requeued = tracker.requeue_worker_videos(wid)
tracker.mark_worker(wid, WS_TIMEOUT, "12hr_timeout")
log.info(f" Requeued {requeued} videos back to pending")
# 2. Poll Kaggle API for each active worker
for w in active:
wid = w["worker_id"]
title = w["kaggle_title"]
if wid in [t["worker_id"] for t in timed_out]:
continue # already handled
kstatus = check_kaggle_kernel_status(title, username)
if not kstatus:
continue
if kstatus == "complete":
log.info(f" βœ… {title} completed")
tracker.mark_worker(wid, WS_COMPLETE, kstatus)
requeued = tracker.requeue_worker_videos(wid, new_status=S_FAILED)
if requeued > 0:
log.info(f" Marked {requeued} silently skipped videos as failed")
elif kstatus == "error":
log.warning(f" ❌ {title} errored")
requeued = tracker.requeue_worker_videos(wid)
tracker.mark_worker(wid, WS_ERROR, kstatus)
log.info(f" Requeued {requeued} videos")
elif kstatus == "cancelled":
log.warning(f" 🚫 {title} cancelled")
requeued = tracker.requeue_worker_videos(wid)
tracker.mark_worker(wid, WS_CANCELLED, kstatus)
log.info(f" Requeued {requeued} videos")
elif kstatus in ("running", "queued"):
tracker.mark_worker(wid, WS_RUNNING, kstatus)
time.sleep(1) # rate limit kaggle API
# ── sync HF status files to update video-level tracking ──────────────
def sync_hf_results(tracker):
hf_token = load_hf_token()
if not hf_token:
return
try:
from huggingface_hub import HfApi, hf_hub_download
api = HfApi(token=hf_token)
files = api.list_repo_files(repo_id=HF_REPO, repo_type="dataset")
status_files = [f for f in files if f.startswith("status/") and f.endswith(".json")]
data_files = [f for f in files if f.startswith("data/") and f.endswith(".parquet")]
for sf in status_files:
try:
local = hf_hub_download(repo_id=HF_REPO, filename=sf, repo_type="dataset",
token=hf_token, force_download=True)
with open(local) as f:
st = json.load(f)
wid = st.get("worker_id","")
state = st.get("state","")
v_status = st.get("video_status", {})
if v_status:
tracker.update_video_statuses(v_status)
log.info(f" HF status: {wid} β†’ done={st.get('done',0)} "
f"failed={st.get('failed',0)} state={state}")
if state == "completed":
tracker.mark_worker(wid, WS_COMPLETE, "hf_confirmed")
except Exception as e:
log.warning(f" Could not read {sf}: {e}")
log.info(f"HF: {len(data_files)} parquet files, {len(status_files)} status files")
except Exception as e:
log.error(f"HF sync error: {e}")
# ── commands ─────────────────────────────────────────────────────────
def cmd_run(args):
"""Autopilot: continuously launch workers, sync, retry β€” respecting Kaggle limits."""
tracker = Tracker()
if not tracker.stats():
tracker.import_from_channels_db()
log.info(f"πŸš€ Autopilot started β€” max {MAX_KAGGLE_SESSIONS} sessions, "
f"{args.videos_per_worker} videos/worker, polling every {POLL_INTERVAL_SEC}s")
log.info(f" Press Ctrl+C to stop\n")
cycle = 0
while True:
cycle += 1
stats = tracker.stats()
pending = stats.get(S_PENDING, 0)
done = stats.get(S_DONE, 0)
total = sum(stats.values())
log.info(f"── Cycle {cycle} ─────────────────────────────────")
log.info(f" Videos: {done}/{total} done, {pending} pending, "
f"{stats.get(S_ASSIGNED,0)} assigned, {stats.get(S_FAILED,0)} failed")
# 1. Sync β€” check which workers finished/failed/timed out
sync_workers(tracker)
# 2. Sync HF results
if cycle % 5 == 0: # every 5 cycles (~10min)
sync_hf_results(tracker)
# 3. Auto-retry failed (every 10 cycles)
if cycle % 10 == 0:
n = tracker.retry_failed(max_retries=3)
if n: log.info(f" Auto-retried {n} failed videos")
# 4. Launch new workers if slots available
slots = tracker.available_slots()
if slots > 0 and pending > 0:
log.info(f" {slots} Kaggle slots free β€” launching workers...")
n = launch_workers(tracker, slots, args.videos_per_worker, push=True)
log.info(f" Launched {n} workers")
elif slots == 0:
log.info(f" All {MAX_KAGGLE_SESSIONS} Kaggle slots in use β€” waiting...")
elif pending == 0:
assigned = stats.get(S_ASSIGNED, 0)
if assigned == 0:
log.info("βœ… All videos processed! Exiting autopilot.")
break
log.info(f" Waiting for {assigned} assigned videos to complete...")
# 5. Sleep
log.info(f" Next check in {POLL_INTERVAL_SEC}s...\n")
try:
time.sleep(POLL_INTERVAL_SEC)
except KeyboardInterrupt:
log.info("\nβ›” Autopilot stopped by user")
break
def cmd_launch(args):
tracker = Tracker()
if not tracker.stats():
tracker.import_from_channels_db()
stats = tracker.stats()
log.info(f"Video stats: {dict(stats)}")
slots = tracker.available_slots()
log.info(f"Kaggle slots available: {slots}/{MAX_KAGGLE_SESSIONS}")
n = min(args.max_workers, slots)
if n <= 0:
log.error(f"No slots! {len(tracker.active_workers())} workers active.")
return
launched = launch_workers(tracker, n, args.videos_per_worker, push=args.push)
log.info(f"Launched {launched} workers" + (" (dry run)" if not args.push else ""))
def cmd_status(args):
tracker = Tracker()
stats = tracker.stats()
total = sum(stats.values()) or 1
active = tracker.active_workers()
slots = tracker.available_slots()
print(f"\n{'='*60}")
print(f" YT Pipeline Status")
print(f"{'='*60}")
print(f" Total videos: {total}")
print(f" Kaggle sessions: {len(active)}/{MAX_KAGGLE_SESSIONS} (slots free: {slots})")
print(f"{'─'*60}")
for s in [S_PENDING, S_ASSIGNED, S_DONE, S_FAILED, S_TIMEOUT]:
c = stats.get(s, 0)
pct = c/total*100
bar = "β–ˆ"*int(pct/2) + "β–‘"*(50-int(pct/2))
print(f" {s:12s} {c:7d} {pct:5.1f}% {bar}")
print(f"{'─'*60}")
workers = tracker.worker_stats()
if workers:
print(f"\n Recent Workers:")
for w in workers:
age = ""
if w["pushed_at"]:
try:
pushed = datetime.fromisoformat(w["pushed_at"])
hrs = (datetime.now(timezone.utc)-pushed).total_seconds()/3600
age = f" ({hrs:.1f}h ago)"
except: pass
ks = '?'
try: ks = w['kaggle_status'] or '?'
except: pass
print(f" {w['worker_id']} {w['kaggle_title']:25s} "
f"v={w['video_count']} {w['status']:10s} "
f"kaggle={ks}{age}")
print()
def cmd_retry(args):
tracker = Tracker()
n = tracker.retry_failed(max_retries=args.max_retries)
print(f"Reset {n} failed/timed-out videos to pending")
def cmd_sync(args):
tracker = Tracker()
sync_workers(tracker)
sync_hf_results(tracker)
def cmd_import(args):
tracker = Tracker()
n = tracker.import_from_channels_db()
print(f"Imported {n} videos")
# ── main ─────────────────────────────────────────────────────────────
def main():
p = argparse.ArgumentParser(description="YT Pipeline Orchestrator v3",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Commands:
run Autopilot β€” continuously launch, sync, retry (respects 10-session limit)
launch One-shot launch (respects slot limit)
status Show pipeline + Kaggle session status
sync Poll Kaggle API + HF for updates
retry Reset failed/timed-out videos for retry
import Import videos from channel_links.db
""")
sub = p.add_subparsers(dest="command")
pr = sub.add_parser("run", help="Autopilot loop")
pr.add_argument("--videos-per-worker", type=int, default=3)
pl = sub.add_parser("launch", help="One-shot launch")
pl.add_argument("--max-workers", type=int, default=5)
pl.add_argument("--videos-per-worker", type=int, default=3)
pl.add_argument("--push", action="store_true")
sub.add_parser("status", help="Show status")
pt = sub.add_parser("retry", help="Retry failed")
pt.add_argument("--max-retries", type=int, default=3)
sub.add_parser("sync", help="Sync from Kaggle+HF")
sub.add_parser("import", help="Import from channel_links.db")
args = p.parse_args()
if not args.command:
p.print_help(); return
{"run":cmd_run,"launch":cmd_launch,"status":cmd_status,
"retry":cmd_retry,"sync":cmd_sync,"import":cmd_import}[args.command](args)
if __name__ == "__main__":
main()