spb/uqo-research-db

1#!/usr/bin/env python32"""Re-balance the download queues (run on M1M32).34Contiguous CIK blocks are badly balanced (old CIKs = large filers).  Companies5already *started* stay on their queue (the downloader resumes them); the6never-started companies are dealt round-robin over all queues so every queue7mixes old and new filers.  Then the downloaders are restarted.8"""9from __future__ import annotations1011import json12import subprocess13import sys14from pathlib import Path1516sys.path.insert(0, str(Path(__file__).resolve().parent))17from orchestrate import BHS64, BHS64_DL, CODE, MASTER, NODES, REMOTE_ROOT, log, ssh  # noqa: E4021819SHARDS = CODE / "data" / "metadata" / "shards.json"202122def listed(host: str, path: str) -> set[str]:23    r = ssh(host, f"ls {path} 2>/dev/null")24    return {l.replace(".json", "").strip() for l in r.stdout.splitlines() if l.strip().endswith(".json")}252627def main() -> None:28    man = json.loads(SHARDS.read_text())29    workers = [n for n in NODES if n != MASTER]30    all_new: set[str] = set()31    for n in workers:32        all_new |= set(man["nodes"][n]["self"]) | set(man["nodes"][n]["bhs64"])3334    started: dict[tuple[str, str], set[str]] = {}35    for n in workers:36        started[(n, "self")] = listed(NODES[n][0], f"{REMOTE_ROOT}/data/metadata/submissions") & all_new37        started[(n, "bhs64")] = listed(BHS64, f"{BHS64_DL}/{n}/submissions") & all_new38    taken = set().union(*started.values())39    unstarted = sorted(all_new - taken)40    log(f"rebalance: {len(all_new)} new CIKs, {len(taken)} started, {len(unstarted)} to redistribute")4142    queues = [(n, k) for n in workers for k in ("self", "bhs64")]43    new_lists = {q: sorted(started[q]) for q in queues}44    for i, cik in enumerate(unstarted):45        new_lists[queues[i % len(queues)]].append(cik)4647    # stop every downloader (files already on disk are kept; the log is append-only)48    for n in workers:49        ssh(NODES[n][0], "pkill -f 'scripts/download_filings.py'; rm -f "50                         f"{REMOTE_ROOT}/data/logs/stage.download_self.done {REMOTE_ROOT}/data/logs/stage.download_self.running")51    ssh(BHS64, "pkill -f 'scripts/download_filings.py'; rm -f " + " ".join(f"{BHS64_DL}/{n}/DONE" for n in workers))5253    shard_dir = CODE / "data" / "metadata" / "shards"54    for (n, k), ciks in new_lists.items():55        (shard_dir / f"{n}.{k}.txt").write_text("\n".join(ciks) + ("\n" if ciks else ""))56        man["nodes"][n][k] = ciks57        man["nodes"][n]["n_total"] = len(man["nodes"][n]["self"]) + len(man["nodes"][n]["bhs64"])58        log(f"  {n:7s} {k:5s} -> {len(ciks):4d} CIKs ({len(started[(n, k)])} resumed)")59    man["rebalanced"] = True60    SHARDS.write_text(json.dumps(man, indent=2))6162    for n in workers:63        subprocess.run(["rsync", "-a", f"{shard_dir}/", f"{NODES[n][0]}:{REMOTE_ROOT}/data/metadata/shards/"], check=True)64    subprocess.run(["rsync", "-a", f"{shard_dir}/", f"{BHS64}:{REMOTE_ROOT}/data/metadata/shards/"], check=True)65    subprocess.run(["rsync", "-a", str(SHARDS), f"{BHS64}:{REMOTE_ROOT}/data/metadata/shards.json"], check=True)66    log("rebalance: lists pushed; relaunching downloads")67    subprocess.run([sys.executable, str(Path(__file__).resolve().parent / "orchestrate.py"), "downloads"], check=True)686970if __name__ == "__main__":71    main()72