SPB Git
spboucher.ai

spb/uqo-research-db

Public

Russell 3000: correctifs du run distribué (section full_document, fusion append-only, décodage Arrow, rééquilibrage des files, clés de rapatriement)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Simon-Pierre Boucher committed (Oct 28, 2026) parent 0304133 Browse files

6 changed files +141 −22

6 changed files
Unified Split
modified scripts/cluster/orchestrate.py +6 −4
@@ -235,16 +235,17 @@ def wait_pipeline(poll: int = 300) -> None:
235 235
236 236
237 def stage_pull() -> None: 237 def stage_pull() -> None:
238− ensure_agent() 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] 239 workers = [n for n in NODES if n != MASTER]
240+ e = "ssh -i ~/.ssh/id_ed25519 -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"
240 241
241 def pull(n: str) -> str: 242 def pull(n: str) -> str:
242 ip = node_host(n) 243 ip = node_host(n)
243 cmd = (f"set -e; cd {REMOTE_ROOT}; mkdir -p data/shards/{n}/export data/shards/{n}/embeddings; " 244 cmd = (f"set -e; cd {REMOTE_ROOT}; mkdir -p data/shards/{n}/export data/shards/{n}/embeddings; "
244− f"rsync -a -e 'ssh -o StrictHostKeyChecking=accept-new' {ip}:{REMOTE_ROOT}/data/export/ data/shards/{n}/export/ && " 245+ f"rsync -a -e '{e}' {ip}:{REMOTE_ROOT}/data/export/ data/shards/{n}/export/ && "
245− f"rsync -a -e 'ssh -o StrictHostKeyChecking=accept-new' {ip}:{REMOTE_ROOT}/data/parquet/embeddings/ data/shards/{n}/embeddings/ && " 246+ f"rsync -a -e '{e}' {ip}:{REMOTE_ROOT}/data/parquet/embeddings/ data/shards/{n}/embeddings/ && "
246 f"test -f data/shards/{n}/export/EXPORT_DONE") 247 f"test -f data/shards/{n}/export/EXPORT_DONE")
247− r = ssh(node_host(MASTER), cmd, agent=True, timeout=6 * 3600) 248+ r = ssh(node_host(MASTER), cmd, timeout=6 * 3600)
248 return f"pull {n}: rc={r.returncode} {r.stderr.strip()[-200:]}" 249 return f"pull {n}: rc={r.returncode} {r.stderr.strip()[-200:]}"
249 250
250 with ThreadPoolExecutor(4) as ex: 251 with ThreadPoolExecutor(4) as ex:
@@ -254,6 +255,7 @@ def stage_pull() -> None:
254 255
255 def stage_merge() -> None: 256 def stage_merge() -> None:
256 m = node_host(MASTER) 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")
257 ssh(m, f"cd {REMOTE_ROOT}; nohup sh -c '.venv/bin/python scripts/merge_shards.py && " 259 ssh(m, f"cd {REMOTE_ROOT}; nohup sh -c '.venv/bin/python scripts/merge_shards.py && "
258 f".venv/bin/python scripts/build_factors.py && touch data/logs/stage.merge.done || touch data/logs/stage.merge.failed' " 260 f".venv/bin/python scripts/build_factors.py && touch data/logs/stage.merge.done || touch data/logs/stage.merge.failed' "
259 f"> data/logs/merge.log 2>&1 < /dev/null &", check=True) 261 f"> data/logs/merge.log 2>&1 < /dev/null &", check=True)
added scripts/cluster/rebalance.py +71 −0
@@ -0,0 +1,71 @@
1+#!/usr/bin/env python3
2+"""Re-balance the download queues (run on M1M32).
3+
4+Contiguous CIK blocks are badly balanced (old CIKs = large filers). Companies
5+already *started* stay on their queue (the downloader resumes them); the
6+never-started companies are dealt round-robin over all queues so every queue
7+mixes old and new filers. Then the downloaders are restarted.
8+"""
9+from __future__ import annotations
10+
11+import json
12+import subprocess
13+import sys
14+from pathlib import Path
15+
16+sys.path.insert(0, str(Path(__file__).resolve().parent))
17+from orchestrate import BHS64, BHS64_DL, CODE, MASTER, NODES, REMOTE_ROOT, log, ssh # noqa: E402
18+
19+SHARDS = CODE / "data" / "metadata" / "shards.json"
20+
21+
22+def 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")}
25+
26+
27+def 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"])
33+
34+ 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_new
37+ started[(n, "bhs64")] = listed(BHS64, f"{BHS64_DL}/{n}/submissions") & all_new
38+ 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")
41+
42+ 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)
46+
47+ # 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))
52+
53+ 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] = ciks
57+ 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"] = True
60+ SHARDS.write_text(json.dumps(man, indent=2))
61+
62+ 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)
68+
69+
70+if __name__ == "__main__":
71+ main()
added scripts/cluster/spid_after_merge.sh +20 −0
@@ -0,0 +1,20 @@
1+#!/bin/bash
2+# Waits for the merged database, then runs the full SPID LLM extraction + post-processing.
3+set -u
4+ROOT="$HOME/UQO_RESEARCH_DB"; cd "$ROOT"
5+export PYTHONUNBUFFERED=1
6+LOGS="$ROOT/data/logs"
7+while [ ! -f "$LOGS/stage.merge.done" ]; do sleep 300; done
8+echo "=== $(date "+%F %T") merge done, starting SPID" > "$LOGS/spid.log"
9+cd project-db
10+PY="$ROOT/.venv/bin/python"
11+PYTHONPATH=. $PY -m spid pending > "$LOGS/spid_pending.json" 2>&1
12+cat "$LOGS/spid_pending.json" >> "$LOGS/spid.log"
13+PYTHONPATH=. $PY run_full_extract.py --workers 24 >> "$LOGS/spid_extract.log" 2>&1
14+rc=$?; echo "=== $(date "+%F %T") extract rc=$rc" >> "$LOGS/spid.log"
15+for step in resolve graph embed analytics; do
16+ echo "=== $(date "+%F %T") $step" >> "$LOGS/spid.log"
17+ PYTHONPATH=. $PY -m spid $step >> "$LOGS/spid_post.log" 2>&1 || echo "$step FAILED" >> "$LOGS/spid.log"
18+done
19+PYTHONPATH=. $PY -m spid stats >> "$LOGS/spid.log" 2>&1
20+touch "$LOGS/stage.spid.done"
modified scripts/merge_shards.py +21 −2
@@ -50,9 +50,19 @@ logger = setup_logging("merge_shards")
50 TABLES = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata"] 50 TABLES = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata"]
51 51
52 52
53+PK = {"filings": "accession_number", "filing_sections": "section_id",
54+ "filing_chunks": "chunk_id", "embeddings_metadata": "embedding_id"}
55+
56+
53 def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None: 57 def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None:
54− for t in TABLES: 58+ """Append the shards' rows that the master does not have yet (keys are globally unique).
59+
60+ ``companies`` is already complete on the master (same universe everywhere); only the
61+ ``shard_node`` tag is updated from the shard's filings.
62+ """
63+ for t in ["filings", "filing_sections", "filing_chunks", "embeddings_metadata"]:
55 cols = [r[0] for r in db.execute(f"DESCRIBE {t}").fetchall()] 64 cols = [r[0] for r in db.execute(f"DESCRIBE {t}").fetchall()]
65+ pk = PK[t]
56 for sd in shard_dirs: 66 for sd in shard_dirs:
57 f = sd / "export" / f"{t}.parquet" 67 f = sd / "export" / f"{t}.parquet"
58 if not f.exists(): 68 if not f.exists():
@@ -61,9 +71,18 @@ def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None:
61 file_cols = [c for c in pq.read_schema(str(f)).names if c in cols] 71 file_cols = [c for c in pq.read_schema(str(f)).names if c in cols]
62 col_list = ", ".join(file_cols) 72 col_list = ", ".join(file_cols)
63 before = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0] 73 before = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]
64− db.execute(f"INSERT OR REPLACE INTO {t} ({col_list}) SELECT {col_list} FROM read_parquet('{f}')") 74+ db.execute(
75+ f"INSERT INTO {t} ({col_list}) SELECT {col_list} FROM read_parquet('{f}') s "
76+ f"WHERE NOT EXISTS (SELECT 1 FROM {t} m WHERE m.{pk} = s.{pk})"
77+ )
65 after = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0] 78 after = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]
66 logger.info("%-20s <- %-8s +%d rows (now %d)", t, sd.name, after - before, after) 79 logger.info("%-20s <- %-8s +%d rows (now %d)", t, sd.name, after - before, after)
80+ for sd in shard_dirs:
81+ f = sd / "export" / "filings.parquet"
82+ if f.exists():
83+ db.execute(f"UPDATE companies SET shard_node = '{sd.name}' WHERE cik IN "
84+ f"(SELECT DISTINCT cik FROM read_parquet('{f}'))")
85+ logger.info("companies.shard_node updated from the shards")
67 86
68 87
69 def move_embeddings(shard_dirs: list[Path]) -> None: 88 def move_embeddings(shard_dirs: list[Path]) -> None:
modified src/factors/scoring.py +10 −9
@@ -36,19 +36,20 @@ logger = get_logger("factors.scoring")
36 36
37 37
38 def _vectors_from_column(col) -> np.ndarray: 38 def _vectors_from_column(col) -> np.ndarray:
39− """Decode a (chunked) list<float|double> Arrow column into a 2-D float32 array without Python loops.""" 39+ """Decode a (chunked) list<float|double> Arrow column into a 2-D float32 array without Python loops.
40+
41+ ``flatten()`` (not ``.values``) is used so that sliced batches with a non-zero offset
42+ decode to exactly ``len(col)`` rows.
43+ """
40 import pyarrow as pa 44 import pyarrow as pa
41 if isinstance(col, pa.ChunkedArray): 45 if isinstance(col, pa.ChunkedArray):
42 col = col.combine_chunks() 46 col = col.combine_chunks()
43− if pa.types.is_fixed_size_list(col.type): 47+ n = len(col)
44− width = col.type.list_size 48+ if n == 0:
45− flat = col.values 49+ return np.zeros((0, 0), dtype=np.float32)
46− else: 50+ flat = col.flatten()
47− offsets = col.offsets.to_numpy()
48− width = int(offsets[1] - offsets[0]) if len(offsets) > 1 else 0
49− flat = col.values
50 arr = np.asarray(flat.to_numpy(zero_copy_only=False), dtype=np.float32) 51 arr = np.asarray(flat.to_numpy(zero_copy_only=False), dtype=np.float32)
51− return arr.reshape(-1, width) if width else arr.reshape(len(col), -1) 52+ return arr.reshape(n, -1)
52 53
53 54
54 class ScoringEngine: 55 class ScoringEngine:
modified src/parsing/filing_parser.py +13 −7
@@ -18,6 +18,7 @@ from __future__ import annotations
18 18
19 import json 19 import json
20 import traceback 20 import traceback
21+import uuid
21 from concurrent.futures import ProcessPoolExecutor, as_completed 22 from concurrent.futures import ProcessPoolExecutor, as_completed
22 from datetime import datetime, timezone 23 from datetime import datetime, timezone
23 from pathlib import Path 24 from pathlib import Path
@@ -78,6 +79,18 @@ class FilingParser:
78 return self._error_result(accession_number, form_type, "empty_after_clean") 79 return self._error_result(accession_number, form_type, "empty_after_clean")
79 80
80 sections = self._extractor.extract_sections(clean_text, form_type) 81 sections = self._extractor.extract_sections(clean_text, form_type)
82+ if not sections and clean_text:
83+ # no Item markers found (e.g. exhibit-only 10-Q/10-K): keep the whole text as one
84+ # real section so the chunks satisfy the filing_chunks -> filing_sections foreign key
85+ sections = [{
86+ "section_id": str(uuid.uuid4()),
87+ "section_name": "full_document",
88+ "section_order": 0,
89+ "item_number": "N/A",
90+ "text": clean_text.strip(),
91+ "word_count": len(clean_text.split()),
92+ "char_count": len(clean_text),
93+ }]
81 for section in sections: 94 for section in sections:
82 section["accession_number"] = accession_number 95 section["accession_number"] = accession_number
83 96
@@ -88,13 +101,6 @@ class FilingParser:
88 chunk["accession_number"] = accession_number 101 chunk["accession_number"] = accession_number
89 chunks.extend(section_chunks) 102 chunks.extend(section_chunks)
90 103
91− if not chunks and clean_text:
92− fallback_chunks = self._chunker.chunk_text(clean_text)
93− for chunk in fallback_chunks:
94− chunk["section_id"] = "full_document"
95− chunk["accession_number"] = accession_number
96− chunks.extend(fallback_chunks)
97−
98 result = { 104 result = {
99 "accession_number": accession_number, 105 "accession_number": accession_number,
100 "form_type": form_type, 106 "form_type": form_type,
101 107
Menu

Theme
↑↓ navigate ↵ open esc close / or ⌘K anywhere