pipeline/export.py

"""Stage 2 — export the OpenSearch indexes and the Neo4j database.

The OpenSearch export deliberately produces the exact file layout that
upstream's `services/deepsearch_api/loader.sh` consumes:

    <index>.settings.json   <index>.mapping.json   <index>.data.json

so the artifacts are restorable by the existing Helm `data_loader_job` and
container-test compose, not only by our own bake step. That compatibility is the
acceptance test for this stage (see the plan, §9 step 5).
"""

from __future__ import annotations

import gzip
import hashlib
import json
import shutil
import subprocess
from pathlib import Path
from typing import Dict, List

import httpx

from .config import BUILD_DIR, REPO_ROOT, Config
from .enumerate_assets import load_manifest, save_manifest

ARTIFACTS = BUILD_DIR / "artifacts"
OS_ARTIFACTS = ARTIFACTS / "opensearch"
NEO4J_ARTIFACTS = ARTIFACTS / "neo4j"

COMPOSE_PROJECT = "simready-indexing"
ELASTICDUMP_IMAGE = "elasticdump/elasticsearch-dump:latest"


def _run(cmd: List[str], **kwargs) -> subprocess.CompletedProcess:
    print("    $ " + " ".join(cmd[:6]) + (" ..." if len(cmd) > 6 else ""))
    return subprocess.run(cmd, check=True, **kwargs)


def _uid_gid() -> str:
    import os
    return f"{os.getuid()}:{os.getgid()}"


def discover_indexes(cfg: Config) -> List[str]:
    """Indexes belonging to this ES_NAME, read from the cluster.

    Discovery rather than a hardcoded pair: the pointcloud cache exists in some
    corpora, and a rename that missed one place would otherwise export an empty
    index without complaint.
    """
    response = httpx.get(f"{cfg.opensearch_url}/_cat/indices?format=json", timeout=30.0)
    response.raise_for_status()
    names = sorted(
        row["index"] for row in response.json()
        if row.get("index", "").startswith(f"{cfg.es_name}-")
    )
    if cfg.main_index not in names:
        raise SystemExit(
            f"main index {cfg.main_index} not found in the cluster. Found: {names or 'none'}. "
            f"Check ES_NAME against what the writers actually created."
        )
    return names


def _network_name() -> str:
    return f"{COMPOSE_PROJECT}_default"


def elasticdump(index: str, dump_type: str, out_path: Path) -> None:
    _run([
        "docker", "run", "--rm",
        "--network", _network_name(),
        "--user", _uid_gid(),
        "-v", f"{OS_ARTIFACTS}:/out",
        ELASTICDUMP_IMAGE,
        f"--input=http://opensearch:9200/{index}",
        f"--output=/out/{out_path.name}",
        f"--type={dump_type}",
        "--limit=1000",
        "--overwrite",
    ])


def sha256_of(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fh:
        for block in iter(lambda: fh.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def gzip_file(path: Path) -> Path:
    dest = path.with_suffix(path.suffix + ".gz")
    with path.open("rb") as src, gzip.open(dest, "wb") as out:
        shutil.copyfileobj(src, out)
    return dest


def export_opensearch(cfg: Config, manifest: Dict) -> Dict:
    OS_ARTIFACTS.mkdir(parents=True, exist_ok=True)
    indexes = discover_indexes(cfg)
    print(f"  indexes: {', '.join(indexes)}")

    expected = manifest.get("asset_count", 0)
    main_count = int(
        httpx.get(f"{cfg.opensearch_url}/{cfg.main_index}/_count", timeout=30.0).json()["count"]
    )
    if expected and main_count != expected:
        raise SystemExit(
            f"pre-flight: {cfg.main_index} holds {main_count} docs but the run expects "
            f"{expected}. Refusing to export a partial index — re-run stage 1."
        )

    # Flush so everything written is on disk and countable.
    httpx.post(f"{cfg.opensearch_url}/_flush", timeout=120.0)

    exported: Dict[str, Dict] = {}
    for index in indexes:
        print(f"  exporting {index} ...")
        files = {}
        for dump_type, suffix in (
            ("settings", "settings.json"),
            ("mapping", "mapping.json"),
            ("data", "data.json"),
        ):
            plain = OS_ARTIFACTS / f"{index}.{suffix}"
            elasticdump(index, dump_type, plain)
            if not plain.exists() or plain.stat().st_size == 0:
                raise SystemExit(f"elasticdump produced no {dump_type} for {index}")
            compressed = gzip_file(plain)
            files[dump_type] = {
                "file": compressed.name,
                "bytes": compressed.stat().st_size,
                "sha256": sha256_of(compressed),
            }
        count = int(
            httpx.get(f"{cfg.opensearch_url}/{index}/_count", timeout=30.0).json()["count"]
        )
        data_lines = sum(
            1 for _ in (OS_ARTIFACTS / f"{index}.data.json").open() if _.strip()
        )
        if data_lines != count:
            raise SystemExit(
                f"{index}: dumped {data_lines} rows but the index reports {count} docs"
            )
        exported[index] = {"doc_count": count, "files": files}
        print(f"    {count} docs")

    return exported


def verify_dump(cfg: Config, dump: Path) -> None:
    """Load the dump into a throwaway data dir to prove it is loadable.

    Same principle as `bake.py::verify_baked` — check the artifact actually
    works from a fresh process, not just that a file got written. A dump can
    fail this while still existing on disk at a plausible size — see the
    stop-ordering comment in `export_neo4j` — and `bake` finding out is
    hours and several stages too late.

    Runs as the host uid, unlike the real dump/load steps: those touch the
    named `neo4j-data` volume, already owned by the in-container uid 7474,
    so they must run as root. This scratch dir is fresh and ours, so running
    as root here would just leave root-owned files behind that the host
    user can't clean up afterwards.
    """
    import tempfile

    with tempfile.TemporaryDirectory() as scratch:
        result = subprocess.run(
            [
                "docker", "run", "--rm",
                "--user", _uid_gid(),
                "-v", f"{dump.parent}:/dumps:ro",
                "-v", f"{scratch}:/data",
                f"neo4j:{cfg.get('NEO4J_VERSION', '5.23.0-community')}",
                "neo4j-admin", "database", "load", "neo4j",
                "--from-path=/dumps", "--overwrite-destination=true",
            ],
            capture_output=True, text=True, check=False,
        )
        if result.returncode != 0:
            raise SystemExit(
                f"neo4j dump failed to verify-load — likely corrupted by AGS "
                f"and neo4j being stopped concurrently (see the stop-ordering "
                f"comment in export_neo4j):\n{result.stdout}\n{result.stderr}"
            )


def export_neo4j(cfg: Config) -> Dict:
    NEO4J_ARTIFACTS.mkdir(parents=True, exist_ok=True)

    # Community edition cannot stop the default database from Cypher (no
    # manual `db.checkpoint()` either — that procedure doesn't exist outside
    # Enterprise), so the dump has to be taken offline via a plain container
    # stop, relying on Neo4j's own graceful-shutdown checkpoint.
    #
    # AGS *must* be given its own `stop` call, fully awaited, before neo4j's
    # `stop` call starts — not one `stop asset-graph-service neo4j` command
    # naming both. `docker compose stop` sends SIGTERM to every named service
    # concurrently regardless of argument order, so a single combined call
    # can (and, reproduced directly, does) stop AGS and neo4j at the same
    # moment — if AGS is mid-write when neo4j's shutdown checkpoint starts,
    # the resulting dump is corrupted (truncated trailing zstd frame,
    # `Not a valid Neo4j archive`) while still landing on disk at a
    # plausible size, silently, caught only by a `neo4j-admin database load`
    # deep in a later `bake`. Two sequential calls make the ordering real
    # instead of advisory; `verify_dump` below is the safety net regardless.
    # `docker compose stop`'s default grace period (~10s) is nowhere near
    # enough for Neo4j's own graceful-shutdown checkpoint on a real corpus —
    # verified live: a ~69k-node/81k-relationship graph took ~103s to stop
    # cleanly. Short of that, compose SIGKILLs it, which produces exactly the
    # "plausible size, silently corrupted" dump the comment above warns
    # about, just from a different cause (timeout, not concurrent-stop
    # ordering) — caught by `verify_dump` below either way, but avoidable.
    print("  stopping asset-graph-service ...")
    _run(["docker", "compose", "-p", COMPOSE_PROJECT, *cfg.indexing_compose_files(),
          "stop", "--timeout", "300", "asset-graph-service"], env=cfg.compose_env())
    # 300 s wasn't enough either on 2026-10-16: the full simready-assets graph
    # (~149k nodes, ~12k property indexes) was SIGKILLed mid-checkpoint and the
    # dump refused with "Active logical log detected". Recover by starting
    # neo4j once (it replays the log) and rerunning export.
    stop_timeout = cfg.get_int("NEO4J_STOP_TIMEOUT", 1800)
    print(f"  stopping neo4j (up to {stop_timeout}s for its shutdown checkpoint) ...")
    _run(["docker", "compose", "-p", COMPOSE_PROJECT, *cfg.indexing_compose_files(),
          "stop", "--timeout", str(stop_timeout), "neo4j"], env=cfg.compose_env())

    print("  dumping neo4j ...")
    # Runs as root (the image default), NOT as the host uid: the data volume is
    # owned by the in-container neo4j user (uid 7474), and neo4j-admin refuses
    # with "You do not have permission to dump the database" under --user.
    # The dump is chowned to the host uid afterwards so the artifact is readable
    # and removable outside the container.
    #
    # rm -f the previous dump ourselves, rather than passing
    # --overwrite-destination=true: that flag has a confirmed upstream bug
    # (neo4j/neo4j#13604) where overwriting a pre-existing *larger* archive
    # doesn't fully truncate it, leaving old bytes at the tail of the new one
    # — `Not a valid Neo4j archive` / `ZstdIOException: Unknown frame
    # descriptor` on load, reproduced here on every retry (each one dumping
    # over the previous attempt's own leftover file). The Neo4j maintainers'
    # own workaround (same issue, 2025-10-27) is to delete the destination
    # first and never pass --overwrite-destination — a genuinely fresh
    # destination file can't hit this bug.
    _run([
        "docker", "run", "--rm",
        "--entrypoint", "bash",
        "-v", f"{COMPOSE_PROJECT}_neo4j-data:/data",
        "-v", f"{NEO4J_ARTIFACTS}:/dumps",
        f"neo4j:{cfg.get('NEO4J_VERSION', '5.23.0-community')}",
        "-c",
        "set -e; rm -f /dumps/neo4j.dump; neo4j-admin database dump neo4j --to-path=/dumps; "
        f"chown -R {_uid_gid()} /dumps",
    ])

    dump = NEO4J_ARTIFACTS / "neo4j.dump"
    if not dump.exists():
        raise SystemExit(f"neo4j-admin produced no dump at {dump}")

    print("  verifying the dump loads cleanly ...")
    verify_dump(cfg, dump)

    compressed = gzip_file(dump)

    print("  restarting neo4j and asset-graph-service ...")
    _run(["docker", "compose", "-p", COMPOSE_PROJECT, *cfg.indexing_compose_files(),
          "start", "neo4j", "asset-graph-service"], env=cfg.compose_env())

    return {
        "file": compressed.name,
        "bytes": compressed.stat().st_size,
        "sha256": sha256_of(compressed),
    }


def package_sizes_problem(manifest: Dict, path: Path) -> str:
    """Why `path` is not this run's package_sizes.json ("" when it is): its (repo, revision) groups must be exactly
    the manifest's hf_repos. On 2026-10-20 a test overwrote .build/package_sizes.json with its fixture
    (nvidia/repo-a@main) and export shipped it with run 20261020-a3cb7f2-captions; this is the check that run
    lacked. Only a storage_api run has one."""
    if manifest.get("storage_backend") != "storage_api":
        return ""
    if not path.exists():
        return f"{path.name} is missing: re-run `enumerate`"
    try:
        groups = json.loads(path.read_text()).get("groups") or []
    except (OSError, ValueError, AttributeError) as e:
        return f"{path.name} is unreadable ({e}): re-run `enumerate`"
    have = {(g.get("repo_id"), g.get("revision")) for g in groups if isinstance(g, dict)}
    want = {(repo_id, revision) for _host, _type, repo_id, revision in manifest.get("hf_repos") or []}
    if have != want:
        extra, missing = sorted(have - want), sorted(want - have)
        return (f"{path.name} is not this run's: {len(extra)} (repo, revision) group(s) the run doesn't list "
                f"{extra[:3]}, {len(missing)} of its own missing {missing[:3]}. Re-run `enumerate` with this "
                f"run's RUN_ID, or restore the file from the run's S3 prefix.")
    return ""


def run(cfg: Config) -> int:
    manifest = load_manifest()
    problem = package_sizes_problem(manifest, BUILD_DIR / "package_sizes.json")
    if problem:
        print(f"export: {problem}")
        return 1
    ARTIFACTS.mkdir(parents=True, exist_ok=True)

    print("Exporting OpenSearch ...")
    manifest["opensearch_artifacts"] = export_opensearch(cfg, manifest)
    from .image_overrides import check_dump  # every doc its own nonzero vector, >= 98 % distinct
    main_dump = OS_ARTIFACTS / f"{cfg.main_index}.data.json.gz"
    if main_dump.exists() and check_dump(main_dump) != 0:
        print(f"export: {main_dump.name} fails the vector check — not exporting a broken index")
        return 1

    # `wait` recorded image_cache_docs before `image-overrides` ran, and each override adds an image document
    # (the package's own thumbnail document stays): 2597 assets -> 2605 image docs on 2026-10-16. scripts/smoke.sh
    # checks the baked image against this number, so take it from what was actually exported.
    image_cache = manifest["opensearch_artifacts"].get(cfg.image_cache_index)
    if image_cache is not None:
        manifest["image_cache_docs"] = image_cache["doc_count"]

    print("\nExporting Neo4j ...")
    manifest["neo4j_artifact"] = export_neo4j(cfg)

    # Copy the inputs alongside the data so an artifact set is self-describing.
    # hf-storage-api-catalog.yaml only exists for a storage_api-backend
    # corpus (pipeline/enumerate_assets.py, stage 0) — folded in here so it
    # rides along with the rest of this run's artifact set through the same
    # upload/download machinery (pipeline/upload.py, pipeline/hf_space_seed.py)
    # rather than needing a separate S3 sync call. See
    # docs/hf-space-deployment.md §4 item 3 for why this is a
    # frozen index-time artifact, not something restored by re-listing the
    # HF dataset at Space startup.
    # package_sizes.json (storage_api too): the Space's /package_sizes.json sums, from enumerate's own listing.
    for name in ("corpus.json", "index-report.json", "assets.jsonl", "hf-storage-api-catalog.yaml",
                 "package_sizes.json"):
        src = BUILD_DIR / name
        if src.exists():
            shutil.copy2(src, ARTIFACTS / name)
        else:
            (ARTIFACTS / name).unlink(missing_ok=True)  # .build/artifacts/ outlives runs; never ship an older one's

    save_manifest(manifest)
    shutil.copy2(BUILD_DIR / "manifest.json", ARTIFACTS / "manifest.json")

    print(f"\nartifacts in {ARTIFACTS.relative_to(REPO_ROOT)}:")
    for path in sorted(ARTIFACTS.rglob("*")):
        if path.is_file():
            print(f"  {path.relative_to(ARTIFACTS)}  ({path.stat().st_size:,} bytes)")
    return 0