diff --git a/CHANGELOG.md b/CHANGELOG.md index f3ed6f7..2b2a99d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ All notable changes to this project are documented in this file. ## Unreleased +### Added +- **`tablassert build-fullmap --aria2c` / `-a`** opt-in downloader acceleration. When requested, the BABEL download stage uses the installed `aria2c` executable with segmented HTTP downloads plus resume/retry flags (`--continue=true`, `--max-tries`, `--retry-wait`) while keeping the existing Python downloader as the default. Missing or failing `aria2c` fails loud instead of silently falling back, and aria2 `.aria2` control files are preserved so interrupted downloads can resume on rerun. + ## 8.2.0 - 2026-08-10 ### Breaking Changes diff --git a/docs/cli.md b/docs/cli.md index 8c2cb8a..03f65d8 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -95,9 +95,12 @@ tablassert build-fullmap [ARGS] | `--cache`, `-c` | Path | No | `./fullmap/downloads` | Directory for downloaded BABEL files (`classes/`, `synonyms/`) | | `--version`, `-v` | str | No | `2026jul22` | BABEL snapshot date to fetch (a RENCI stamp, **not** Tablassert's version) | | `--threads`, `-t` | int | No | `None` (auto) | Worker threads; auto-capped by memory on Linux (`/proc/meminfo`), else ~90% of CPUs | +| `--aria2c`, `-a` | Flag | No | `False` | Opt into the installed `aria2c` executable for resumable segmented BABEL downloads; fails loud if `aria2c` is missing or exits non-zero | ```bash tablassert build-fullmap --output /data/fullmap/fullmap.redb +# Optional: use aria2c for faster/resumable BABEL downloads when installed +tablassert build-fullmap --aria2c --output /data/fullmap/fullmap.redb ``` See [Fullmap](fullmap.md) for the data pipeline, output schema, and graph-config usage. diff --git a/docs/fullmap.md b/docs/fullmap.md index 0b96410..e2297de 100644 --- a/docs/fullmap.md +++ b/docs/fullmap.md @@ -10,18 +10,21 @@ files (`fullmap.s0.redb` … `fullmap.s15.redb` by default) holding the term→p biological synonyms, CURIEs, Biolink categories, taxon IDs, and source provenance, built from NCATS Translator BABEL export files. -Fullmap is built entirely in-process by Tablassert's own Rust extension — no external tool or install step required (this is an in-process redb shard scheme, not the older external DuckDB shards). +Fullmap is built entirely in-process by Tablassert's own Rust extension — no external tool or install step required by default (this is an in-process redb shard scheme, not the older external DuckDB shards). If you opt into `build-fullmap --aria2c` / `-a`, only the download stage uses an installed external `aria2c` executable. ## Build Command ```bash # Build a fullmap database (downloads BABEL data automatically) tablassert build-fullmap + +# Optional: use installed aria2c for resumable segmented BABEL downloads +tablassert build-fullmap --aria2c ``` See the [CLI Reference → build-fullmap](cli.md#build-fullmap) for the complete flag table (output path, -cache directory, BABEL snapshot version, worker threads), their defaults, and more examples. Two facts -matter most when planning a build: +cache directory, BABEL snapshot version, worker threads, and the optional `--aria2c` / `-a` downloader), +their defaults, and more examples. Two facts matter most when planning a build: - The BABEL **version** flag selects a RENCI BABEL snapshot date (default `2026jul22`) — *not* Tablassert's package version. Bumping it fetches a different snapshot and requires rebuilding; the @@ -34,7 +37,7 @@ matter most when planning a build: The build is a parallel, **memory-bounded** pipeline executed by the Rust extension: -1. **Download** — fetch BABEL class and synonym files from RENCI into the cache (resumable, reused). +1. **Download** — fetch BABEL class and synonym files from RENCI into the cache (resumable, reused). By default this uses Tablassert's Python downloader; `--aria2c` / `-a` opts into the installed `aria2c` executable, preserving aria2 resume control files across dropped downloads and failing loud if the executable is missing or the download fails. 2. **Equivalents index** — parse class files into sorted on-disk runs, then k-way merge them into a memory-mapped index mapping each primary CURIE to its equivalents. 3. **Synonym pass** — a producer/consumer pool streams byte-bounded line-chunks; workers dedup CURIEs, @@ -50,7 +53,10 @@ The build is a parallel, **memory-bounded** pipeline executed by the Rust extens multi-threaded allocation from bloating resident memory. - **Download** — files come from `https://stars.renci.org/var/babel_outputs` via resumable, - range-request downloads; cached files are reused. + range-request downloads; cached files are reused. Passing `--aria2c` / `-a` switches only this + stage to the installed `aria2c` executable, using aria2's segmented HTTP downloads and retry/resume + control files while suppressing aria2's own progress UI so Tablassert's progress bar stays clean. + The progress detail remains file-level (`aria2c downloading`) rather than byte-level in this mode. - **Equivalents index** — class files parse in parallel into sorted on-disk runs, k-way merged into a single memory-mapped CURIE→equivalents index; only a compact `(hash, offset)` index lives in RAM, the string data is mmap'd. diff --git a/src/tablassert/cli.py b/src/tablassert/cli.py index 3baae3f..4024b15 100644 --- a/src/tablassert/cli.py +++ b/src/tablassert/cli.py @@ -1,6 +1,8 @@ from __future__ import annotations import re +import shutil +import subprocess import sys import time from collections.abc import Callable @@ -420,6 +422,93 @@ def download_babel_file(filename: str, url: str, destination: Path, retries: int raise BabelDownloadError(url, retries, last_error or RuntimeError("no attempts made")) from last_error +def download_babel_file_aria2c(filename: str, url: str, destination: Path, retries: int = 5) -> Path: + """Download one BABEL file with the optional external ``aria2c`` executable. + + The helper mirrors ``download_babel_file``'s final-file cache contract but + delegates resume/retry behavior to aria2. Incomplete aria2 downloads leave a + ``.aria2`` control file next to the target; when that control file + exists we do NOT treat the target as a cache hit, and failures never remove + either file so a later run can continue. + + Args: + filename: Output basename under ``destination``. + url: Source URL. + destination: Directory to download into (created if missing). + retries: Maximum aria2 tries (forwarded to ``--max-tries``). + + Returns: + Path to the downloaded file. + + Raises: + BabelDownloadError: If ``aria2c`` is missing, fails, or does not leave a + complete final file. + """ + destination.mkdir(parents=True, exist_ok=True) + final_path: Path = destination / filename + control_path: Path = destination / f"{filename}.aria2" + if final_path.is_file() and not control_path.exists(): + download_logger.info("Reusing cached BABEL file: {path}", path=final_path) + return final_path + if retries < 1: + error = ValueError("aria2c retries must be a positive integer") + raise BabelDownloadError(url, retries, error) from error + + binary: str | None = shutil.which("aria2c") + if binary is None: + error = FileNotFoundError("aria2c executable not found; install aria2 or omit --aria2c") + raise BabelDownloadError(url, 0, error) from error + + command: list[str] = [ + binary, + "--continue=true", + "--max-tries", + str(retries), + "--retry-wait", + "5", + "--allow-overwrite=true", + "--auto-file-renaming=false", + "--max-connection-per-server=8", + "--split=8", + "--min-split-size=1M", + "--summary-interval=0", + "--console-log-level=warn", + "--show-console-readout=false", + "--dir", + str(destination), + "--out", + filename, + url, + ] + try: + completed: subprocess.CompletedProcess[str] = subprocess.run( + command, shell=False, check=False, capture_output=True, text=True, errors="replace" + ) + except OSError as e: + raise BabelDownloadError(url, retries, e) from e + + if completed.returncode != 0: + output: str = (completed.stderr or completed.stdout or "").strip() + detail: str = f"aria2c exited with status {completed.returncode}" + if output: + detail = f"{detail}: {output[-2000:]}" + error = RuntimeError(detail) + raise BabelDownloadError(url, retries, error) from error + + if not final_path.is_file() or control_path.exists(): + suffix: str = "" + if control_path.exists(): + suffix = f"; resume control file still present: {control_path}" + output = (completed.stderr or completed.stdout or "").strip() + if output: + suffix = f"{suffix}; aria2c output: {output[-2000:]}" + error = FileNotFoundError(f"aria2c completed but did not create a complete file at {final_path}{suffix}") + raise BabelDownloadError(url, retries, error) from error + + download_logger.info("Downloaded {url} -> {path} with aria2c", url=url, path=final_path) + return final_path + + def stream_copy(source: BinaryIO, destination: BinaryIO, on_bytes: Callable[[int], None] | None = None) -> None: """Copy ``source`` to ``destination`` in 1 MiB chunks. @@ -780,7 +869,12 @@ def rebuild_agent_graph( def build_fullmap_pipeline( - output: Path, progress: PipelineProgress, cache: Path = Path("./fullmap/downloads"), version: str = BABEL_VERSION, threads: int | None = None + output: Path, + progress: PipelineProgress, + cache: Path = Path("./fullmap/downloads"), + version: str = BABEL_VERSION, + threads: int | None = None, + aria2c: bool = False, ) -> None: """Build an embedded fullmap redb database from BABEL outputs. @@ -793,6 +887,7 @@ def build_fullmap_pipeline( cache: Directory for downloaded BABEL files. version: BABEL version label. threads: Optional thread count forwarded to Rust. + aria2c: Use the optional aria2c executable for downloads when true. """ from tablassert import rs @@ -816,18 +911,19 @@ def build_fullmap_pipeline( def report_progress(downloaded: int, total: int) -> None: sub_step(_download_detail(downloaded, total)) - class_files: list[Path] = [] - for filename, url in class_urls: + def download_one(filename: str, url: str, destination: Path) -> Path: start(filename) - sub_step("downloading") - class_files.append(download_babel_file(filename, url, cache / "classes", on_progress=report_progress)) - advance() - synonym_files: list[Path] = [] - for filename, url in synonym_urls: - start(filename) - sub_step("downloading") - synonym_files.append(download_babel_file(filename, url, cache / "synonyms", on_progress=report_progress)) + if aria2c: + sub_step("aria2c downloading") + path: Path = download_babel_file_aria2c(filename, url, destination) + else: + sub_step("downloading") + path = download_babel_file(filename, url, destination, on_progress=report_progress) advance() + return path + + class_files: list[Path] = [download_one(filename, url, cache / "classes") for filename, url in class_urls] + synonym_files: list[Path] = [download_one(filename, url, cache / "synonyms") for filename, url in synonym_urls] # Stage 3/3: build fullmap database. progress.stage("Building Fullmap Database") @@ -852,6 +948,7 @@ def build_fullmap( cache: Annotated[Path, cyclopts.Parameter(name=["--cache", "-c"])] = Path("./fullmap/downloads"), version: Annotated[str, cyclopts.Parameter(name=["--version", "-v"])] = BABEL_VERSION, threads: Annotated[int | None, cyclopts.Parameter(name=["--threads", "-t"])] = None, + aria2c: Annotated[bool, cyclopts.Parameter(name=["--aria2c", "-a"], negative="")] = False, ) -> None: """Build an embedded fullmap redb database from hardcoded BABEL outputs.""" - run(3, build_fullmap_pipeline, output, cache=cache, version=version, threads=threads) + run(3, build_fullmap_pipeline, output, cache=cache, version=version, threads=threads, aria2c=aria2c) diff --git a/tests/test_cover_cli.py b/tests/test_cover_cli.py index 6d9ecc9..6f14495 100644 --- a/tests/test_cover_cli.py +++ b/tests/test_cover_cli.py @@ -10,6 +10,7 @@ from __future__ import annotations import io +import subprocess from email.message import Message from pathlib import Path from typing import Any @@ -19,7 +20,7 @@ from cyclopts.exceptions import UnknownOptionError # pyright: ignore[reportMissingImports] from tablassert import cli, rs -from tablassert.cli import build_fullmap_pipeline, build_kg, download_babel_file, validate_graph_pipeline +from tablassert.cli import build_fullmap_pipeline, build_kg, download_babel_file, download_babel_file_aria2c, validate_graph_pipeline from tablassert.errors import BabelDownloadError, GraphValidationError from tablassert.ingests import to_yaml from tablassert.progress import PipelineProgress @@ -184,6 +185,128 @@ def test_download_babel_file_zero_retries_raises_without_attempt(tmp_path: Path) download_babel_file("f.gz", "https://example.com/f.gz", tmp_path, retries=0) +def test_download_babel_file_aria2c_reuses_cached_complete_without_binary(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A complete final file is reused before binary detection or subprocess execution.""" + final: Path = tmp_path / "f.gz" + final.write_bytes(b"cached-bytes") + + def _which_must_not_run(name: str) -> str | None: + raise AssertionError(f"shutil.which({name!r}) must not run for a complete cache hit") + + def _run_must_not_run(*args: Any, **kwargs: Any) -> subprocess.CompletedProcess[str]: + raise AssertionError("subprocess.run must not run for a complete cache hit") + + monkeypatch.setattr(cli.shutil, "which", _which_must_not_run) + monkeypatch.setattr(cli.subprocess, "run", _run_must_not_run) + out: Path = download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path) + assert out == final + assert out.read_bytes() == b"cached-bytes" + + +def test_download_babel_file_aria2c_runs_resume_retry_command(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """aria2c helper uses subprocess without shell and passes resume/retry flags.""" + monkeypatch.setattr(cli.shutil, "which", lambda name: "/usr/bin/aria2c" if name == "aria2c" else None) + commands: list[list[str]] = [] + + def _fake_run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + commands.append(command) + assert kwargs["shell"] is False + destination = Path(command[command.index("--dir") + 1]) + filename = command[command.index("--out") + 1] + destination.mkdir(parents=True, exist_ok=True) + (destination / filename).write_bytes(b"downloaded") + return subprocess.CompletedProcess(command, 0, stdout="", stderr="") + + monkeypatch.setattr(cli.subprocess, "run", _fake_run) + out: Path = download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path, retries=7) + assert out == tmp_path / "f.gz" + assert out.read_bytes() == b"downloaded" + assert len(commands) == 1 + command = commands[0] + assert command[0] == "/usr/bin/aria2c" + assert "--continue=true" in command + assert "--max-tries" in command + assert command[command.index("--max-tries") + 1] == "7" + assert "--retry-wait" in command + assert command[command.index("--retry-wait") + 1] == "5" + assert "--summary-interval=0" in command + assert "--show-console-readout=false" in command + assert command[-1] == "https://example.com/f.gz" + + +def test_download_babel_file_aria2c_missing_binary_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Opting into aria2c fails loud when the executable is unavailable.""" + monkeypatch.setattr(cli.shutil, "which", lambda name: None) + + def _run_must_not_run(*args: Any, **kwargs: Any) -> subprocess.CompletedProcess[str]: + raise AssertionError("subprocess.run must not run when aria2c is missing") + + monkeypatch.setattr(cli.subprocess, "run", _run_must_not_run) + with pytest.raises(BabelDownloadError): + download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path) + + +def test_download_babel_file_aria2c_zero_retries_raises_without_unlimited_aria2(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """``retries=0`` is rejected before aria2 can interpret it as unlimited retries.""" + + def _which_must_not_run(name: str) -> str | None: + raise AssertionError("aria2c lookup must not run when retries is invalid") + + monkeypatch.setattr(cli.shutil, "which", _which_must_not_run) + with pytest.raises(BabelDownloadError): + download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path, retries=0) + + +def test_download_babel_file_aria2c_subprocess_oserror_raises_typed_error(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """OS errors from launching aria2c surface as ``BabelDownloadError``.""" + monkeypatch.setattr(cli.shutil, "which", lambda name: "/usr/bin/aria2c") + + def _raise_oserror(*args: Any, **kwargs: Any) -> subprocess.CompletedProcess[str]: + raise OSError("exec failed") + + monkeypatch.setattr(cli.subprocess, "run", _raise_oserror) + with pytest.raises(BabelDownloadError): + download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path) + + +def test_download_babel_file_aria2c_success_without_complete_file_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Even an exit-0 aria2c run must leave a final file without a resume control file.""" + monkeypatch.setattr(cli.shutil, "which", lambda name: "/usr/bin/aria2c") + + def _fake_run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + destination = Path(command[command.index("--dir") + 1]) + filename = command[command.index("--out") + 1] + destination.mkdir(parents=True, exist_ok=True) + (destination / filename).write_bytes(b"partial") + (destination / f"{filename}.aria2").write_bytes(b"resume-state") + return subprocess.CompletedProcess(command, 0, stdout="done", stderr="") + + monkeypatch.setattr(cli.subprocess, "run", _fake_run) + with pytest.raises(BabelDownloadError): + download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path) + + +def test_download_babel_file_aria2c_preserves_control_file_on_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """An existing aria2 control file means the target is partial and must be resumed, not reused.""" + final: Path = tmp_path / "f.gz" + control: Path = tmp_path / "f.gz.aria2" + final.write_bytes(b"partial") + control.write_bytes(b"resume-state") + monkeypatch.setattr(cli.shutil, "which", lambda name: "/usr/bin/aria2c") + commands: list[list[str]] = [] + + def _fake_run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]: + commands.append(command) + return subprocess.CompletedProcess(command, 1, stdout="", stderr="dropped connection") + + monkeypatch.setattr(cli.subprocess, "run", _fake_run) + with pytest.raises(BabelDownloadError): + download_babel_file_aria2c("f.gz", "https://example.com/f.gz", tmp_path) + assert len(commands) == 1 # final + .aria2 was NOT treated as a complete cache hit + assert final.read_bytes() == b"partial" + assert control.read_bytes() == b"resume-state" # failure path preserves aria2 resume metadata + + def test_build_kg_command_delegates_to_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """Cover cli.py:501 — the ``build-kg`` cyclopts command forwards to ``run(6, build_pipeline, ...)``. @@ -202,6 +325,35 @@ def _fake_run(stages: int, fn: Any, arg: Path, **kwargs: Any) -> None: assert calls == [(6, cli.build_pipeline, config, {"release": True, "qc": True, "log": True, "head": True})] +def test_build_fullmap_command_passes_aria2c_flag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """``build-fullmap --aria2c`` delegates the opt-in flag to ``build_fullmap_pipeline``.""" + output: Path = tmp_path / "fullmap.redb" + cache: Path = tmp_path / "downloads" + calls: list[tuple[Any, ...]] = [] + + def _fake_run(stages: int, fn: Any, arg: Path, **kwargs: Any) -> None: + calls.append((stages, fn, arg, kwargs)) + + monkeypatch.setattr(cli, "run", _fake_run) + cli.build_fullmap(output=output, cache=cache, version="v", threads=2, aria2c=True) + assert calls == [(3, cli.build_fullmap_pipeline, output, {"cache": cache, "version": "v", "threads": 2, "aria2c": True})] + + +def test_build_fullmap_aria2c_flag_parses() -> None: + """``build-fullmap`` accepts ``--aria2c`` and ``-a`` but no generated negative alias.""" + + def parse(argv: list[str]) -> dict[str, Any]: + fn, bound, _ = cli.APP.parse_args(argv, exit_on_error=False) + assert fn is cli.build_fullmap + return dict(bound.arguments) + + assert parse(["build-fullmap"]) == {} + assert parse(["build-fullmap", "--aria2c"])["aria2c"] is True + assert parse(["build-fullmap", "-a"])["aria2c"] is True + with pytest.raises(UnknownOptionError): + parse(["build-fullmap", "--no-aria2c"]) + + def test_build_kg_configuration_file_flag_parses(tmp_path: Path) -> None: """Guard: ``build-kg``'s config binds positionally AND via ``-f``/``--configuration-file``. @@ -277,3 +429,71 @@ def _fake_build(output: Path, class_files: list[Path], synonym_files: list[Path] assert (cache / "synonyms" / "s.gz").read_bytes() == payload # Stage 3 received the downloaded paths and the thread count. assert built == [(output, [cache / "classes" / "c.gz"], [cache / "synonyms" / "s.gz"], 1)] + + +def test_build_fullmap_pipeline_uses_aria2c_when_opted_in(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The aria2c flag switches both class and synonym loops to the aria2 helper and simple progress text.""" + + def _fake_babel_urls(version: str, endpoints: tuple[str, ...], pattern: object) -> list[tuple[str, str]]: + if endpoints == cli.BABEL_CLASS_ENDPOINTS: + return [("c.gz", "https://example.com/c.gz")] + return [("s.gz", "https://example.com/s.gz")] + + monkeypatch.setattr(cli, "babel_urls", _fake_babel_urls) + + def _python_download_must_not_run(*args: Any, **kwargs: Any) -> Path: + raise AssertionError("Python downloader must not run when aria2c=True") + + monkeypatch.setattr(cli, "download_babel_file", _python_download_must_not_run) + aria_calls: list[tuple[str, str, Path]] = [] + + def _fake_aria2c(filename: str, url: str, destination: Path, retries: int = 5) -> Path: + aria_calls.append((filename, url, destination)) + destination.mkdir(parents=True, exist_ok=True) + path = destination / filename + path.write_bytes(b"downloaded") + return path + + monkeypatch.setattr(cli, "download_babel_file_aria2c", _fake_aria2c) + built: list[tuple[Any, ...]] = [] + + def _fake_build(output: Path, class_files: list[Path], synonym_files: list[Path], threads: int | None = None, progress: Any = None) -> None: + built.append((output, class_files, synonym_files, threads)) + + monkeypatch.setattr(rs, "build_fullmap_db", _fake_build) + + class _RecordingProgress: + def __init__(self) -> None: + self.sub_steps: list[str] = [] + self.advances: int = 0 + + def stage(self, name: str) -> None: + pass + + def section_loop(self, total: int, label: str) -> tuple[Any, Any, Any]: + def start(detail: str) -> None: + pass + + def advance() -> None: + self.advances += 1 + + def sub_step(phase: str) -> None: + self.sub_steps.append(phase) + + return start, advance, sub_step + + def dynamic_loop(self, label: str) -> Any: + return lambda *args: None + + def end_section_task(self) -> None: + pass + + progress = _RecordingProgress() + output: Path = tmp_path / "fullmap.redb" + cache: Path = tmp_path / "downloads" + build_fullmap_pipeline(output, progress, cache=cache, version="v", threads=1, aria2c=True) # type: ignore[arg-type] + + assert aria_calls == [("c.gz", "https://example.com/c.gz", cache / "classes"), ("s.gz", "https://example.com/s.gz", cache / "synonyms")] + assert progress.sub_steps.count("aria2c downloading") == 2 + assert progress.advances == 4 # two discovery entries + two downloaded files + assert built == [(output, [cache / "classes" / "c.gz"], [cache / "synonyms" / "s.gz"], 1)]