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>
6 changed files +141 −22
modified
scripts/cluster/orchestrate.py
+6 −4
@@ -235,16 +235,17 @@ def wait_pipeline(poll: int = 300) -> None: | ||
| 235 | 235 | |
| 236 | 236 | |
| 237 | 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 | 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 | 242 | def pull(n: str) -> str: |
| 242 | 243 | ip = node_host(n) |
| 243 | 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 'ssh -o StrictHostKeyChecking=accept-new' {ip}:{REMOTE_ROOT}/data/parquet/embeddings/ 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/ && " | |
| 246 | 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 | 249 | return f"pull {n}: rc={r.returncode} {r.stderr.strip()[-200:]}" |
| 249 | 250 | |
| 250 | 251 | with ThreadPoolExecutor(4) as ex: |
@@ -254,6 +255,7 @@ def stage_pull() -> None: | ||
| 254 | 255 | |
| 255 | 256 | def stage_merge() -> None: |
| 256 | 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 | 259 | ssh(m, f"cd {REMOTE_ROOT}; nohup sh -c '.venv/bin/python scripts/merge_shards.py && " |
| 258 | 260 | f".venv/bin/python scripts/build_factors.py && touch data/logs/stage.merge.done || touch data/logs/stage.merge.failed' " |
| 259 | 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 | 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 | 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 | 64 | cols = [r[0] for r in db.execute(f"DESCRIBE {t}").fetchall()] |
| 65 | + pk = PK[t] | |
| 56 | 66 | for sd in shard_dirs: |
| 57 | 67 | f = sd / "export" / f"{t}.parquet" |
| 58 | 68 | if not f.exists(): |
@@ -61,9 +71,18 @@ def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None: | ||
| 61 | 71 | file_cols = [c for c in pq.read_schema(str(f)).names if c in cols] |
| 62 | 72 | col_list = ", ".join(file_cols) |
| 63 | 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 | 78 | after = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0] |
| 66 | 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 | 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 | 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 | 44 | import pyarrow as pa |
| 41 | 45 | if isinstance(col, pa.ChunkedArray): |
| 42 | 46 | col = col.combine_chunks() |
| 43 | − if pa.types.is_fixed_size_list(col.type): | |
| 44 | − width = col.type.list_size | |
| 45 | − flat = col.values | |
| 46 | − else: | |
| 47 | − offsets = col.offsets.to_numpy() | |
| 48 | − width = int(offsets[1] - offsets[0]) if len(offsets) > 1 else 0 | |
| 49 | − flat = col.values | |
| 47 | + n = len(col) | |
| 48 | + if n == 0: | |
| 49 | + return np.zeros((0, 0), dtype=np.float32) | |
| 50 | + flat = col.flatten() | |
| 50 | 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 | 55 | class ScoringEngine: |
modified
src/parsing/filing_parser.py
+13 −7
@@ -18,6 +18,7 @@ from __future__ import annotations | ||
| 18 | 18 | |
| 19 | 19 | import json |
| 20 | 20 | import traceback |
| 21 | +import uuid | |
| 21 | 22 | from concurrent.futures import ProcessPoolExecutor, as_completed |
| 22 | 23 | from datetime import datetime, timezone |
| 23 | 24 | from pathlib import Path |
@@ -78,6 +79,18 @@ class FilingParser: | ||
| 78 | 79 | return self._error_result(accession_number, form_type, "empty_after_clean") |
| 79 | 80 | |
| 80 | 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 | 94 | for section in sections: |
| 82 | 95 | section["accession_number"] = accession_number |
| 83 | 96 | |
@@ -88,13 +101,6 @@ class FilingParser: | ||
| 88 | 101 | chunk["accession_number"] = accession_number |
| 89 | 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 | 104 | result = { |
| 99 | 105 | "accession_number": accession_number, |
| 100 | 106 | "form_type": form_type, |
| 101 | 107 | |