Spaces:
Running on RTX PRO 6000
Running on RTX PRO 6000
| """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 | |