spb/uqo-research-db
Public
1#!/usr/bin/env python32"""Cluster orchestrator for the Russell 3000 extension (runs on the M1M32 gateway).34 python3 orchestrate.py sync | downloads | ship | pipeline | pull | merge | status | all56State: every stage is idempotent and polls marker files on the nodes, so the7script can be stopped and restarted at any stage. Logs: ~/edgar-orchestrate/.8"""910from __future__ import annotations1112import json13import os14import subprocess15import sys16import time17from concurrent.futures import ThreadPoolExecutor18from datetime import datetime19from pathlib import Path2021HOME = Path.home()22CODE = HOME / "UQO_RESEARCH_DB-code" # code + metadata mirror on the gateway23WORK = HOME / "edgar-orchestrate"24WORK.mkdir(exist_ok=True)25LOG = WORK / "orchestrate.log"2627MASTER = "M3U96b"28BHS64 = "ubuntu@51.161.112.61"29BHS64_DL = "edgar-dl" # ~/edgar-dl/<node>/{raw,submissions,download_log.jsonl,DONE}30REMOTE_ROOT = "UQO_RESEARCH_DB"3132# node: (lan_ip, parse_workers)33NODES = {34 "M3U96b": ("192.168.2.85", 26),35 "M3U96a": ("192.168.2.80", 24),36 "M2U64": ("192.168.2.61", 20),37 "M5M36": ("192.168.2.104", 14),38 "M4M64a": ("192.168.2.70", 12),39 "M4M64c": ("192.168.2.103", 14),40 "M4BP48": ("192.168.2.114", 12),41 "M4BP36": ("192.168.2.113", 10),42 "m4mm": ("192.168.2.92", 10),43 "m4mn": ("192.168.2.93", 10),44}45LAN_RATE_TOTAL = 9.0 # SEC: 10 req/s per IP; all LAN nodes share one public IP46BHS64_RATE_TOTAL = 9.04748SSH_OPTS = ["-o", "BatchMode=yes", "-o", "ConnectTimeout=20", "-o", "StrictHostKeyChecking=accept-new",49 "-o", "ServerAliveInterval=30"]505152def log(msg: str) -> None:53 line = f"{datetime.now():%F %T} {msg}"54 print(line, flush=True)55 with open(LOG, "a") as f:56 f.write(line + "\n")575859def ssh(host: str, cmd: str, agent: bool = False, timeout: int = 600, check: bool = False) -> subprocess.CompletedProcess:60 args = ["ssh", "-n", *SSH_OPTS] + (["-A"] if agent else []) + [host, cmd]61 r = subprocess.run(args, capture_output=True, text=True, timeout=timeout)62 if check and r.returncode != 0:63 raise RuntimeError(f"ssh {host} failed rc={r.returncode}: {r.stderr.strip()[:500]}")64 return r656667def node_host(node: str) -> str:68 return NODES[node][0]697071def ensure_agent() -> None:72 """Start an ssh-agent holding the gateway key so nodes can rsync from BHS64 / the master from nodes."""73 if os.environ.get("SSH_AUTH_SOCK") and subprocess.run(["ssh-add", "-l"], capture_output=True).returncode == 0:74 return75 out = subprocess.run(["ssh-agent", "-s"], capture_output=True, text=True, check=True).stdout76 for part in out.replace("\n", ";").split(";"):77 part = part.strip()78 if part.startswith("SSH_AUTH_SOCK=") or part.startswith("SSH_AGENT_PID="):79 k, v = part.split("=", 1)80 os.environ[k] = v81 subprocess.run(["ssh-add", str(HOME / ".ssh" / "id_ed25519")], check=True, capture_output=True)82 log("ssh-agent started with id_ed25519")838485# ----------------------------------------------------------------------------- stages86def stage_sync() -> None:87 excl = ["--exclude", ".DS_Store", "--exclude", "__pycache__", "--exclude", ".venv", "--exclude", "data/",88 "--exclude", "project-db/", "--exclude", "analysis/", "--exclude", "notebooks/",89 "--exclude", "outputs/", "--exclude", "UQ_WP_template/", "--exclude", "*.zip", "--exclude", ".git/"]90 meta = ["universe.json", "shards.json", "russell3000_constituents.json", "sp500_constituents.json", "cik_mapping.json"]9192 def push(target: str) -> str:93 subprocess.run(["rsync", "-a", *excl, "-e", "ssh " + " ".join(SSH_OPTS), f"{CODE}/", f"{target}:{REMOTE_ROOT}/"],94 check=True, capture_output=True)95 subprocess.run(["ssh", *SSH_OPTS, target, f"mkdir -p {REMOTE_ROOT}/data/metadata/shards {REMOTE_ROOT}/data/logs"],96 check=True, capture_output=True)97 subprocess.run(["rsync", "-a", "-e", "ssh " + " ".join(SSH_OPTS),98 *[str(CODE / "data" / "metadata" / m) for m in meta if (CODE / "data" / "metadata" / m).exists()],99 f"{target}:{REMOTE_ROOT}/data/metadata/"], check=True, capture_output=True)100 subprocess.run(["rsync", "-a", "-e", "ssh " + " ".join(SSH_OPTS), f"{CODE}/data/metadata/shards/",101 f"{target}:{REMOTE_ROOT}/data/metadata/shards/"], check=True, capture_output=True)102 return target103104 with ThreadPoolExecutor(8) as ex:105 for t in ex.map(push, [node_host(n) for n in NODES] + [BHS64]):106 log(f"sync -> {t} ok")107108109def stage_downloads() -> None:110 shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]111 lan_nodes = [n for n in NODES if shards.get(n, {}).get("self")]112 rate = LAN_RATE_TOTAL / max(1, len(lan_nodes))113 for n in lan_nodes:114 r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.download_self.done 2>/dev/null; "115 f"pgrep -f 'shards/[{n[0]}]{n[1:]}.self.txt' >/dev/null && echo RUNNING")116 if "done" in r.stdout or "RUNNING" in r.stdout:117 log(f"downloads: {n} already {'done' if 'done' in r.stdout else 'running'}")118 continue119 ssh(node_host(n), f"cd {REMOTE_ROOT}; RATE={rate:.3f} DL_WORKERS=6 nohup scripts/cluster/run_shard.sh {n} download_self "120 f"> data/logs/nohup.download_self.out 2>&1 < /dev/null &", check=True)121 log(f"downloads: {n} started self list ({len(shards[n]['self'])} CIKs @ {rate:.2f} req/s)")122123 bhs_nodes = [n for n in NODES if shards.get(n, {}).get("bhs64")]124 brate = BHS64_RATE_TOTAL / max(1, len(bhs_nodes))125 for n in bhs_nodes:126 r = ssh(BHS64, f"cat {BHS64_DL}/{n}/DONE 2>/dev/null; pgrep -f 'shards/[{n[0]}]{n[1:]}.bhs64.txt' >/dev/null && echo RUNNING")127 if "rc=" in r.stdout or "RUNNING" in r.stdout:128 log(f"downloads: BHS64/{n} already {'done' if 'rc=' in r.stdout else 'running'}")129 continue130 cmd = (f"mkdir -p {BHS64_DL}/{n}; cd {REMOTE_ROOT}; PYTHONUNBUFFERED=1 nohup sh -c "131 f"'.venv/bin/python scripts/download_filings.py --ciks-file data/metadata/shards/{n}.bhs64.txt "132 f"--rate {brate:.3f} --workers 6 --raw-dir ~/{BHS64_DL}/{n}/raw --log ~/{BHS64_DL}/{n}/download_log.jsonl "133 f"--submissions-dir ~/{BHS64_DL}/{n}/submissions; echo rc=$? > ~/{BHS64_DL}/{n}/DONE' "134 f"> ~/{BHS64_DL}/{n}/download.log 2>&1 < /dev/null &")135 ssh(BHS64, cmd, check=True)136 log(f"downloads: BHS64 started for {n} ({len(shards[n]['bhs64'])} CIKs @ {brate:.2f} req/s)")137138139def downloads_status() -> dict:140 shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]141 st = {}142 for n in NODES:143 r = ssh(node_host(n), f"cd {REMOTE_ROOT}/data/logs 2>/dev/null && ls stage.download_self.* 2>/dev/null; "144 f"tail -1 {REMOTE_ROOT}/data/logs/download_self.log 2>/dev/null | cut -c1-160")145 st[n] = r.stdout.strip().replace("\n", " | ")146 for n in NODES:147 if not shards.get(n, {}).get("bhs64"):148 continue149 r = ssh(BHS64, f"cat {BHS64_DL}/{n}/DONE 2>/dev/null; echo \"$(wc -l < {BHS64_DL}/{n}/download_log.jsonl 2>/dev/null) filings\"; tail -1 {BHS64_DL}/{n}/download.log 2>/dev/null | cut -c1-160")150 st[f"BHS64/{n}"] = r.stdout.strip().replace("\n", " | ")151 return st152153154def wait_downloads(poll: int = 300) -> None:155 shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]156 while True:157 pending = []158 for n in NODES:159 if shards.get(n, {}).get("self"):160 r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.download_self.done {REMOTE_ROOT}/data/logs/stage.download_self.failed 2>/dev/null")161 if "done" not in r.stdout and "failed" not in r.stdout:162 pending.append(n)163 elif "failed" in r.stdout:164 log(f"WARNING {n} download_self FAILED (partial data kept; rerun 'downloads' later)")165 if shards.get(n, {}).get("bhs64"):166 r = ssh(BHS64, f"cat {BHS64_DL}/{n}/DONE 2>/dev/null")167 if "rc=" not in r.stdout:168 pending.append(f"BHS64/{n}")169 if not pending:170 log("downloads: all finished")171 return172 log(f"downloads: waiting on {pending}")173 time.sleep(poll)174175176def stage_ship() -> None:177 """Move the BHS64-downloaded halves onto their nodes (node pulls over SSH with the gateway agent)."""178 ensure_agent()179 shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]180 targets = [n for n in NODES if shards.get(n, {}).get("bhs64")]181182 def pull(n: str) -> str:183 r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.ship.done 2>/dev/null")184 if r.returncode == 0:185 return f"{n}: already shipped"186 cmd = (f"set -e; cd {REMOTE_ROOT}; mkdir -p data/raw/sec data/metadata/submissions; "187 f"rsync -a --partial -e 'ssh -o StrictHostKeyChecking=accept-new' {BHS64}:{BHS64_DL}/{n}/raw/ data/raw/sec/ && "188 f"rsync -a -e 'ssh -o StrictHostKeyChecking=accept-new' {BHS64}:{BHS64_DL}/{n}/submissions/ data/metadata/submissions/ && "189 f"rsync -a -e 'ssh -o StrictHostKeyChecking=accept-new' {BHS64}:{BHS64_DL}/{n}/download_log.jsonl data/metadata/download_log.bhs64.jsonl && "190 f"touch data/logs/stage.ship.done")191 r = ssh(node_host(n), cmd, agent=True, timeout=6 * 3600)192 return f"{n}: ship rc={r.returncode} {r.stderr.strip()[-200:]}"193194 with ThreadPoolExecutor(len(targets) or 1) as ex:195 for msg in ex.map(pull, targets):196 log(msg)197198199def stage_pipeline() -> None:200 while True: # the master needs the S&P 500 corpus (rsync'ed from the laptop) before it can parse201 r = ssh(node_host(MASTER), f"ls {REMOTE_ROOT}/data/logs/stage.master_corpus.done 2>/dev/null")202 if r.returncode == 0:203 break204 log("pipeline: waiting for stage.master_corpus.done on the master (S&P 500 corpus transfer)")205 time.sleep(300)206 for n in NODES:207 r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.pipeline.done {REMOTE_ROOT}/data/logs/stage.pipeline.running 2>/dev/null")208 if "done" in r.stdout or "running" in r.stdout:209 log(f"pipeline: {n} already {'done' if 'done' in r.stdout else 'running'}")210 continue211 workers = NODES[n][1]212 stage = "pipeline_master" if n == MASTER else "pipeline"213 ssh(node_host(n), f"cd {REMOTE_ROOT}; WORKERS={workers} nohup scripts/cluster/run_shard.sh {n} {stage} "214 f"> data/logs/nohup.pipeline.out 2>&1 < /dev/null &", check=True)215 log(f"pipeline: {n} started ({workers} parse workers)")216217218def wait_pipeline(poll: int = 300) -> None:219 while True:220 pending, failed = [], []221 for n in NODES:222 r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/ 2>/dev/null | grep '^stage\\.' | tr '\\n' ' '")223 marks = r.stdout224 if "stage.pipeline.done" in marks:225 continue226 if "stage.pipeline.failed" in marks:227 failed.append(n)228 else:229 pending.append(n)230 if not pending:231 log(f"pipeline: finished (failed: {failed})")232 return233 log(f"pipeline: waiting on {pending} failed={failed}")234 time.sleep(poll)235236237def stage_pull() -> None:238 """The master pulls every worker's export + embeddings with its own key (authorised on the workers)."""239 workers = [n for n in NODES if n != MASTER]240 e = "ssh -i ~/.ssh/id_ed25519 -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"241242 def pull(n: str) -> str:243 ip = node_host(n)244 cmd = (f"set -e; cd {REMOTE_ROOT}; mkdir -p data/shards/{n}/export data/shards/{n}/embeddings; "245 f"rsync -a -e '{e}' {ip}:{REMOTE_ROOT}/data/export/ data/shards/{n}/export/ && "246 f"rsync -a -e '{e}' {ip}:{REMOTE_ROOT}/data/parquet/embeddings/ data/shards/{n}/embeddings/ && "247 f"test -f data/shards/{n}/export/EXPORT_DONE")248 r = ssh(node_host(MASTER), cmd, timeout=6 * 3600)249 return f"pull {n}: rc={r.returncode} {r.stderr.strip()[-200:]}"250251 with ThreadPoolExecutor(4) as ex:252 for msg in ex.map(pull, workers):253 log(msg)254255256def stage_merge() -> None:257 m = node_host(MASTER)258 ssh(m, f"rm -f {REMOTE_ROOT}/data/logs/stage.merge.done {REMOTE_ROOT}/data/logs/stage.merge.failed")259 ssh(m, f"cd {REMOTE_ROOT}; nohup sh -c '.venv/bin/python scripts/merge_shards.py && "260 f".venv/bin/python scripts/build_factors.py && touch data/logs/stage.merge.done || touch data/logs/stage.merge.failed' "261 f"> data/logs/merge.log 2>&1 < /dev/null &", check=True)262 log("merge: started on master (merge_shards + build_factors)")263 while True:264 r = ssh(m, f"ls {REMOTE_ROOT}/data/logs/stage.merge.* 2>/dev/null; tail -1 {REMOTE_ROOT}/data/logs/merge.log | cut -c1-200")265 if "merge.done" in r.stdout or "merge.failed" in r.stdout:266 log(f"merge: {r.stdout.strip()}")267 return268 log(f"merge: running | {r.stdout.strip().splitlines()[-1] if r.stdout.strip() else ''}")269 time.sleep(300)270271272def stage_status() -> None:273 for n, (ip, _) in NODES.items():274 r = ssh(ip, f"cd {REMOTE_ROOT}/data/logs 2>/dev/null && ls | grep '^stage\\.' | tr '\\n' ' '; "275 f"df -h / | tail -1 | awk '{{print $4\" free\"}}'; "276 f"echo \"self-log: $(cat ~/{REMOTE_ROOT}/data/metadata/download_log.self.jsonl 2>/dev/null | wc -l | tr -d ' ') filings\"; "277 f"for f in download_self parse embed export merge; do [ -f {REMOTE_ROOT}/data/logs/$f.log ] && tail -1 {REMOTE_ROOT}/data/logs/$f.log | cut -c1-150; done")278 print(f"--- {n}\n{r.stdout.strip()}")279 for k, v in downloads_status().items():280 if k.startswith("BHS64"):281 print(f"--- {k}: {v}")282283284STAGES = {285 "sync": stage_sync, "downloads": stage_downloads, "wait-downloads": wait_downloads, "ship": stage_ship,286 "pipeline": stage_pipeline, "wait-pipeline": wait_pipeline, "pull": stage_pull, "merge": stage_merge,287 "status": stage_status,288}289290291def main() -> None:292 if len(sys.argv) < 2 or sys.argv[1] not in STAGES and sys.argv[1] != "all":293 print(__doc__)294 sys.exit(1)295 what = sys.argv[1]296 if what == "all":297 for s in ["sync", "downloads", "wait-downloads", "ship", "pipeline", "wait-pipeline", "pull", "merge"]:298 log(f"##### stage {s}")299 STAGES[s]()300 log("##### ALL DONE")301 else:302 STAGES[what]()303304305if __name__ == "__main__":306 main()307