f4092e5331
За неделю обе массовые заливки умерли молча и не возобновились: Википедия на 20 008 статьях из ~2 млн («server closed the connection»), PMC примерно на 85 тыс. из 100 тыс. (таймаут SSL-рукопожатия). Ни ретраев, ни возобновления. - запись пачки повторяется с переподключением к базе (5 попыток с паузой); - обрыв листинга бакета стоит паузы, а не всей заливки; - у Википедии сохраняется позиция в дампе — после сбоя продолжаем с неё, а не проматываем с нуля, полагаясь на дедуп по ext_id; - батч уменьшен (500 статей × 2000 отпечатков = миллион строк в одной транзакции — вероятная причина обрыва); - обрезка отпечатков переведена на равномерную выборку. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
248 lines
12 KiB
Python
Executable File
248 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""Массовая заливка полнотекстовых статей из открытого бакета PMC (миллионы).
|
||
|
||
Почему отдельный путь, а не парсеры и Celery: постраничный API даёт ~1-2 статьи
|
||
в секунду и упирается в rate limit — миллионы так не залить. У AWS Open Data
|
||
бакета `pmc-oa-opendata` (открыт, ключ не нужен) для каждой статьи лежит уже
|
||
извлечённый текст `PMC*.txt` (33-130 КБ) и метаданные `PMC*.json`. Замер:
|
||
12 статей/с в 10 потоков, то есть миллион — примерно за сутки.
|
||
|
||
Второе узкое место — запись отпечатков. Построчный INSERT даёт ~6.7 тыс. строк/с
|
||
(миллион статей = 2.5 млрд строк = четверо суток только на вставку), поэтому
|
||
здесь используется COPY: он на порядки быстрее и не гоняет данные через ORM.
|
||
|
||
Масштаб для планирования: 1 млн статей ≈ 2.5 млрд отпечатков ≈ 400 ГБ в БД
|
||
с индексами. Плотность отпечатков ограничена --fp-per-doc (по умолчанию 2000
|
||
вместо 20 000 из настроек воркера) — иначе объём растёт быстрее пользы.
|
||
|
||
Запуск (в контейнере worker-indexer, где есть psycopg2 и winnowing):
|
||
cd /home/user/anti-plagiarism
|
||
C="docker compose -f docker-compose.prod.yml exec -T worker-indexer python -"
|
||
$C --limit 200 < scripts/ops/bulk_ingest_pmc.py # пробная порция
|
||
$C --limit 100000 --workers 20 < scripts/ops/bulk_ingest_pmc.py
|
||
|
||
Идемпотентно: документы вставляются с ON CONFLICT DO NOTHING по ext_id, уже
|
||
залитые пропускаются. Прогресс листинга сохраняется в --state-file, поэтому
|
||
прерванная заливка продолжается с той же страницы бакета.
|
||
|
||
Проверено на проде: 1831 отпечаток на статью — реальная глубина, а не аннотация.
|
||
Темп ~2 статьи/с в один поток (0.3 ст/с из первого замера были получены, когда
|
||
база одновременно считала эмбеддинги). Скачивание не узкое место (12 статей/с в
|
||
12 потоков), упирается в последовательную обработку документа — для миллионов
|
||
её нужно распараллелить по ядрам воркера.
|
||
"""
|
||
|
||
import argparse
|
||
import contextlib
|
||
import io
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from urllib.parse import quote
|
||
|
||
BUCKET = "https://pmc-oa-opendata.s3.amazonaws.com"
|
||
STATE_FILE = "/tmp/bulk_pmc_state.json"
|
||
KEY_RE = re.compile(r"<Key>(PMC\d+\.\d+)/\1\.txt</Key>")
|
||
TOKEN_RE = re.compile(r"<NextContinuationToken>([^<]+)</NextContinuationToken>")
|
||
|
||
|
||
def list_articles(client, token: str | None, want: int) -> tuple[list[str], str | None]:
|
||
"""Собрать id статей из листинга бакета, продолжая с сохранённой страницы."""
|
||
ids: list[str] = []
|
||
while len(ids) < want:
|
||
url = f"{BUCKET}/?list-type=2&max-keys=1000"
|
||
if token:
|
||
# В токене бывают + и / — без экранирования S3 отвечает 400
|
||
url += f"&continuation-token={quote(token, safe='')}"
|
||
resp = client.get(url)
|
||
resp.raise_for_status()
|
||
ids.extend(KEY_RE.findall(resp.text))
|
||
m = TOKEN_RE.search(resp.text)
|
||
token = m.group(1) if m else None
|
||
if token is None:
|
||
break
|
||
return ids[:want], token
|
||
|
||
|
||
def fetch_one(client, art_id: str) -> dict | None:
|
||
"""Текст статьи и её метаданные. None — если статья недоступна или пустая."""
|
||
try:
|
||
txt = client.get(f"{BUCKET}/{art_id}/{art_id}.txt")
|
||
if txt.status_code != 200 or len(txt.text) < 1500:
|
||
return None
|
||
meta_resp = client.get(f"{BUCKET}/{art_id}/{art_id}.json")
|
||
meta = meta_resp.json() if meta_resp.status_code == 200 else {}
|
||
except Exception:
|
||
return None
|
||
|
||
citation = meta.get("citation") or ""
|
||
year = None
|
||
m = re.search(r"\b(19|20)\d{2}\b", citation)
|
||
if m:
|
||
year = int(m.group(0))
|
||
|
||
return {
|
||
"ext_id": f"pmc:{meta.get('pmcid') or art_id.split('.')[0]}",
|
||
"title": (meta.get("title") or art_id)[:1000],
|
||
"doi": meta.get("doi"),
|
||
"year": year,
|
||
"journal": citation[:300] or None,
|
||
"text": txt.text,
|
||
}
|
||
|
||
|
||
|
||
|
||
def write_batch(conn, connect, docs: list[dict], fp_limit: int) -> tuple:
|
||
"""Записать пачку статей, пережив обрыв соединения с базой.
|
||
|
||
Прошлые прогоны умирали целиком от `server closed the connection`, теряя и
|
||
позицию в бакете. Здесь обрыв стоит переподключения и повтора.
|
||
|
||
Returns:
|
||
(соединение, добавлено, пропущено, отпечатков) — соединение может быть
|
||
новым, если пришлось переподключаться
|
||
"""
|
||
import psycopg2
|
||
from app.algorithms.winnowing import sample_evenly, winnow_ordered
|
||
|
||
for attempt in range(1, 6):
|
||
try:
|
||
added = skipped = fp_count = 0
|
||
with conn.cursor() as cur:
|
||
# ON CONFLICT — повторный прогон не плодит дубли
|
||
cur.executemany(
|
||
"""INSERT INTO documents (ext_id, source, title, doi, year, lang, journal,
|
||
abstract, authors, indexed_at)
|
||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,'[]'::json, now())
|
||
ON CONFLICT (ext_id) DO NOTHING""",
|
||
[(d["ext_id"], "pmc", d["title"], d["doi"], d["year"], "en",
|
||
d["journal"], d["text"][:2000]) for d in docs],
|
||
)
|
||
cur.execute("SELECT id, ext_id FROM documents WHERE ext_id = ANY(%s)",
|
||
([d["ext_id"] for d in docs],))
|
||
id_by_ext = {e: i for i, e in cur.fetchall()}
|
||
|
||
# COPY вместо построчных INSERT — единственный способ писать
|
||
# миллиарды строк за разумное время
|
||
buf = io.StringIO()
|
||
for d in docs:
|
||
doc_id = id_by_ext.get(d["ext_id"])
|
||
if doc_id is None:
|
||
continue
|
||
cur.execute("SELECT 1 FROM fingerprints WHERE doc_id=%s LIMIT 1", (doc_id,))
|
||
if cur.fetchone():
|
||
skipped += 1
|
||
continue
|
||
# Равномерно по тексту, а не первые N: иначе хвост статьи
|
||
# остаётся без отпечатков
|
||
hashes = sample_evenly(winnow_ordered(d["text"]), fp_limit)
|
||
for pos, h in enumerate(hashes):
|
||
buf.write(f"{doc_id}\t{h}\t{pos}\n")
|
||
fp_count += len(hashes)
|
||
added += 1
|
||
buf.seek(0)
|
||
if added:
|
||
cur.copy_from(buf, "fingerprints",
|
||
columns=("doc_id", "hash_value", "position"))
|
||
conn.commit()
|
||
return conn, added, skipped, fp_count
|
||
|
||
except psycopg2.OperationalError as e:
|
||
wait = min(60, 5 * attempt)
|
||
print(f" база недоступна ({str(e).strip()[:60]}), повтор через {wait}с "
|
||
f"[{attempt}/5]", flush=True)
|
||
time.sleep(wait)
|
||
with contextlib.suppress(Exception):
|
||
conn.close()
|
||
with contextlib.suppress(Exception):
|
||
conn = connect()
|
||
|
||
raise RuntimeError("не удалось записать пачку после 5 попыток")
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description=__doc__,
|
||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
ap.add_argument("--limit", type=int, default=200, help="сколько статей залить за прогон")
|
||
ap.add_argument("--workers", type=int, default=10, help="потоков закачки (вежливо: 10-30)")
|
||
ap.add_argument("--batch", type=int, default=200, help="статей в одной транзакции")
|
||
ap.add_argument("--fp-per-doc", type=int, default=2000, help="максимум отпечатков на документ")
|
||
ap.add_argument("--state-file", default=STATE_FILE, help="файл с позицией листинга")
|
||
args = ap.parse_args()
|
||
|
||
import httpx
|
||
import psycopg2
|
||
from app.config import settings
|
||
|
||
token = None
|
||
if os.path.exists(args.state_file):
|
||
with open(args.state_file) as fh:
|
||
token = json.load(fh).get("token")
|
||
print("продолжаю листинг с сохранённой позиции")
|
||
|
||
client = httpx.Client(timeout=60,
|
||
headers={"User-Agent": "AcademicHelper/1.0 (noreply@jze9.ru)"})
|
||
|
||
def connect():
|
||
return psycopg2.connect(
|
||
host=settings.POSTGRES_HOST, port=settings.POSTGRES_PORT,
|
||
dbname=settings.POSTGRES_DB, user=settings.POSTGRES_USER,
|
||
password=settings.POSTGRES_PASSWORD,
|
||
)
|
||
|
||
conn = connect()
|
||
added = skipped = failed = fp_total = done = 0
|
||
t0 = time.time()
|
||
|
||
while done < args.limit:
|
||
want = min(args.batch, args.limit - done)
|
||
# Сеть у бакета иногда отваливается (прошлый прогон умер на таймауте
|
||
# SSL-рукопожатия) — это стоит паузы, а не всей заливки
|
||
try:
|
||
ids, token = list_articles(client, token, want)
|
||
except Exception as e:
|
||
print(f" листинг не удался ({type(e).__name__}), пауза 30с", flush=True)
|
||
time.sleep(30)
|
||
continue
|
||
|
||
if not ids:
|
||
print("бакет закончился")
|
||
break
|
||
done += len(ids)
|
||
|
||
with ThreadPoolExecutor(max_workers=args.workers) as pool:
|
||
docs = [d for d in pool.map(lambda i: fetch_one(client, i), ids) if d]
|
||
failed += len(ids) - len(docs)
|
||
if not docs:
|
||
continue
|
||
|
||
conn, n_added, n_skipped, n_fp = write_batch(conn, connect, docs, args.fp_per_doc)
|
||
added += n_added
|
||
skipped += n_skipped
|
||
fp_total += n_fp
|
||
|
||
with open(args.state_file, "w") as fh:
|
||
json.dump({"token": token}, fh)
|
||
|
||
el = time.time() - t0
|
||
print(f" {done}/{args.limit} · залито {added} · пропущено {skipped} · "
|
||
f"недоступно {failed} · отпечатков {fp_total:,} · {added / max(el, 1):.1f} ст/с",
|
||
flush=True)
|
||
if token is None:
|
||
break
|
||
|
||
conn.close()
|
||
el = time.time() - t0
|
||
print(f"\nГотово за {el / 60:.1f} мин: залито {added}, пропущено {skipped}, "
|
||
f"недоступно {failed}, отпечатков {fp_total:,}")
|
||
if added:
|
||
print(f"темп: {added / el:.1f} статей/с → "
|
||
f"миллион за ~{1e6 / max(added / el, 0.01) / 3600:.0f} ч")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|