diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b02ba77..5a581bd 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -7,6 +7,34 @@ on: branches: ["**"] jobs: + dependency-resolution: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Resolve dependencies without GitHub credentials + env: + GIT_TERMINAL_PROMPT: "0" + run: | + python -m venv /tmp/pathbench-clean + /tmp/pathbench-clean/bin/python -m pip install --upgrade pip + /tmp/pathbench-clean/bin/python -m pip install \ + --dry-run --ignore-installed --report /tmp/install-report.json . + /tmp/pathbench-clean/bin/python - <<'PY' + import json + from pathlib import Path + + report = json.loads(Path("/tmp/install-report.json").read_text()) + urls = [item["download_info"]["url"] for item in report["install"]] + vcs_urls = [url for url in urls if url.startswith("git+") or "github.com" in url] + assert not vcs_urls, f"VCS/GitHub dependencies remain: {vcs_urls}" + PY + test: runs-on: ubuntu-latest strategy: diff --git a/README.md b/README.md index 2382dc2..e9d6f09 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,25 @@ score = evaluator.score("utt1", "/path/to/audio.wav", transcription="the cat sat print(f"ArtP score: {score}") ``` +ArtP is reference-based: it force-aligns the phonemes in the supplied transcription +with the audio, so the transcription and its language are required. DArtP is +reference-free: a language-specific ASR model and n-gram language model first +produce a transcription, which is then scored by the same phonetic model: + +```python +from pathbench import ArtPDoubleASREvaluator + +evaluator = ArtPDoubleASREvaluator(language="en-us") +score = evaluator.score("utt1", "/path/to/audio.wav") +print(f"DArtP score: {score}") +``` + +DArtP currently supports `en`/`en-us`, `es`, `nl`, `it`, and `cmn`. It also +requires the corresponding file from the [n-gram model download](#n-gram-models) +to be placed in `lms/`; without it, scoring returns `None`. Run commands from +the repository root because DArtP resolves `lms/` relative to the working +directory. + ### I want to contribute a new predictor to this repository, how do I do that? See [CONTRIBUTING.md](CONTRIBUTING.md) for a step-by-step guide. @@ -120,18 +139,117 @@ Results are written to the `results_11/` directory as timestamped text files con ## Installation -We are continously trying to make the installation easier for your use case. +We are continuously trying to make the installation easier for your use case. + +### GPU/Colab native bootstrap + +The helper makes **no host changes by default**. Without either installation flag it +only validates the pinned native backend, NVIDIA visibility, and the Python/GPU +environment: + +```bash +python3.12 tools/test_gpu_predictors.py --python python3.12 --cuda-version 12.4 --pytorch-version 2.6.0 +``` + +`--install-system-dependencies` (or +`PATHBENCH_INSTALL_SYSTEM_DEPENDENCIES=1`; an explicit command-line setting takes +precedence) is an opt-in to run noninteractive `apt-get update` and install `git`, +`ca-certificates`, `curl`, `build-essential`, `cmake`, `ninja-build`, `pkg-config`, +`libfftw3-dev`, and `liblapack-dev`. It adds `python3-venv` only for a distribution +Python that lacks `venv`, not for a uv-managed interpreter. On Debian/Ubuntu it +then builds espeak-ng at pinned commit **`2ea41210`**, installs its executable and +shared library under `/usr/local`, refreshes the linker cache, and records success +at `/usr/local/share/pathbench/espeak-ng-commit`. Root or passwordless +noninteractive sudo is required. Other operating systems are rejected with the +manual prerequisite list. A failed build/probe never writes the marker. + +#### Google Colab + +Colab owns the NVIDIA driver: this helper never replaces it, changes kernel +modules/CUDA libraries, or reboots Colab. Select a GPU runtime first; if +`nvidia-smi --list-gpus` fails, select a GPU runtime or reconnect to a new runtime. +Acquire Python 3.12 using uv and clone the repository outside the helper, then run: + +```bash +# after installing uv, Python 3.12, and cloning/cd'ing into PathBench +uv python install 3.12 +git clone https://github.com/karkirowle/pathbench.git && cd pathbench +python3.12 tools/test_gpu_predictors.py --python python3.12 --cuda-version 12.4 --pytorch-version 2.6.0 --install-system-dependencies --download-language-model +``` + +A fresh T4 validation should install the build packages, build pinned espeak-ng, +create the CUDA environment, verify the English model, and finish with `2 passed`. +A second invocation should report reuse of the marker, virtual environment, +CUDA packages, dependencies, and verified model, again ending with `2 passed`. +Colab storage is ephemeral. The driver-install option is always refused there. + +#### Container + +Native user-space dependencies may be installed inside a Debian/Ubuntu image: + +```bash +python tools/test_gpu_predictors.py --install-system-dependencies --download-language-model +``` + +Configure the NVIDIA driver and NVIDIA Container Toolkit/runtime **on the host**, +then expose the GPU to the container. Installing a host kernel driver inside a +container, Kubernetes pod, WSL, or other GPU-passthrough guest is refused; +restarting that guest cannot activate a host kernel module. + +#### Bare-metal Ubuntu + +Bootstrap ordinary dependencies independently of the driver: + +```bash +python tools/test_gpu_predictors.py --install-system-dependencies --download-language-model +``` + +Driver handling is separate and experimental, intended only for an +administrator-controlled bare-metal Ubuntu host. The first command is a +non-destructive preflight: it reports `lspci`, the kernel and matching headers, +and the signed package recommended by `ubuntu-drivers` (never inferred from the +CUDA wheel): + +```bash +python tools/test_gpu_predictors.py --install-nvidia-driver +# Review the output, then explicitly approve package changes: +python tools/test_gpu_predictors.py --install-nvidia-driver --confirm-nvidia-driver-install +sudo reboot +# Only after reboot, rerun normally; nvidia-smi --list-gpus must succeed. +python tools/test_gpu_predictors.py --install-system-dependencies --download-language-model +``` + +The confirmed option runs apt and installs the Ubuntu-recommended +`nvidia-driver-*` package, then stops with a successful **reboot required** +message without creating a venv or running tests. It never uses NVIDIA's `.run` +installer. `PATHBENCH_INSTALL_NVIDIA_DRIVER=1` can request preflight, but the +confirmation intentionally has no environment equivalent. + +**Maintainer bare-metal integration checklist (do not run driver replacement in +CI):** use a disposable administrator-controlled Ubuntu host; record `lspci`, +`uname -r`, Secure Boot state, current `nvidia-smi`, and recommended package; run +the unconfirmed preflight and verify it changes nothing; review apt's proposed +changes; snapshot the host; confirm manually; verify the reboot-required exit; +reboot from the host console; then verify `nvidia-smi --list-gpus` before running +the helper again. Never perform this checklist in Colab, a container, WSL, a +Kubernetes pod, or a provider-managed GPU VM. If you have the opportunity to start from a clean AWS/GCE instance, please do so and follow the make installation. If you are working on a highly restricted HPC cluster, I would recommend starting from the singularity container [provided](https://github.com/karkirowle/pathbench/releases/download/v0.1.0/pathbench.sif). -Package installation is the recommended pathway when you are trying to incorporate into your existing stuff. In this case, you are kind of your own figuring out +Package installation is the recommended pathway when incorporating PathBench +into an existing environment. In that case, you are responsible for resolving dependency conflicts. -### Package installation +All Python runtime dependencies are available from package indexes rather than +VCS URLs. In particular, `phonemizer-fork==3.3.2` (which installs the +`phonemizer` import package) and `pyctcdecode==0.5.0` use versioned PyPI +releases, so installing PathBench does not require GitHub credentials. The +system-level espeak-ng revision below remains separately pinned because its +language-specific IPA output is part of the metric definition. -PathBench cannot be published to PyPI because it depends on Git-hosted forks of `phonemizer` and `pyctcdecode`. +### Package installation **System dependencies** (not installable via pip — must be installed separately): - `espeak-ng` at commit [`2ea41210`](https://github.com/espeak-ng/espeak-ng/commit/2ea41210) (post-1.52.0) — required by the phonemizer for grapheme-to-phoneme conversion. The exact commit matters: different espeak-ng versions produce different IPA symbols for some languages (e.g. Italian `ɾ` vs `r`), which affects phoneme-based metrics (PER, dPER, ArtP). Build from source: @@ -141,7 +259,8 @@ PathBench cannot be published to PyPI because it depends on Git-hosted forks of cmake -B build -DUSE_ASYNC=OFF -DBUILD_SHARED_LIBS=ON cmake --build build -j$(nproc) && sudo cmake --install build ``` -- PyTorch with CUDA support — install following [pytorch.org](https://pytorch.org/get-started/locally/) *before* installing pathbench +- PyTorch — install the CPU or CUDA build appropriate for your system by following + [pytorch.org](https://pytorch.org/get-started/locally/) *before* installing PathBench. **Option A — Install from a GitHub Release:** ```bash @@ -160,7 +279,10 @@ pip install "pathbench[scripts] @ git+https://github.com/karkirowle/pathbench.gi ### Make installation -The `make` installation route assumes the default setup of a standard Ubuntu 22.04 image (`ubuntu-2204-jammy`). +The `make` installation route assumes the default setup of a standard Ubuntu +22.04 image (`ubuntu-2204-jammy`) and Python 3.10–3.12. It creates +`tools/venv`. The default is a CPU-only PyTorch installation, which works for +inference and tests but is slower than a supported GPU. ```bash sudo apt-get update -qq @@ -171,12 +293,120 @@ cd /tmp/espeak-ng && git checkout 2ea41210 cmake -B build -DUSE_ASYNC=OFF -DBUILD_SHARED_LIBS=ON cmake --build build -j$(nproc) && sudo cmake --install build && sudo ldconfig cd - -git clone git@github.com:karkirowle/pathbench.git +git clone https://github.com/karkirowle/pathbench.git cd pathbench/tools && make cd .. source tools/venv/bin/activate ``` +For a CUDA build, select a wheel index supported by the pinned PyTorch version. +For example, PyTorch 2.6.0 provides CUDA 12.4 wheels: + +```bash +cd pathbench/tools +make CUDA_VERSION=12.4 +``` + +You can select a particular interpreter with, for example, +`make PYTHON=python3.12`. Re-running `make` resumes after completed stages; +run `make clean` first to rebuild the environment with a different Python, +PyTorch, or CUDA selection. + +### GPU installation and predictor smoke test + +After installing the pinned `espeak-ng` build above, systems with an NVIDIA GPU +and driver can use the helper script to create a separate CUDA environment and +run the focused ArtP and DArtP tests: + +```bash +python tools/test_gpu_predictors.py --download-language-model --cuda-version 12.4 +``` + +The script checks for `nvidia-smi` and `espeak-ng`, creates +`tools/gpu_venv`, installs the CUDA 12.4 builds of PyTorch and torchaudio 2.6.0, +installs PathBench and its test dependencies, verifies that PyTorch can access +the GPU, and runs both predictor tests. It can be invoked from any directory. +The first run downloads the Python packages and model checkpoints and therefore +requires network access and several gigabytes of free disk space. + +Override its defaults with command-line options (or the corresponding +`PYTHON`, `PATHBENCH_CUDA_VERSION`, `PYTORCH_VERSION`, and `VENV` environment variables) +when needed. The selected CUDA wheel must exist for the selected PyTorch release: + +```bash +python tools/test_gpu_predictors.py --python python3.11 --cuda-version 12.6 \ + --pytorch-version 2.6.0 --venv /path/to/pathbench-gpu-venv +``` + +The script requires an NVIDIA driver compatible with the chosen CUDA wheel; +installing the wheel does not install a host GPU driver or the CUDA toolkit. +Language-model download is deliberately opt-in. With +`--download-language-model`, the helper uses HTTP range requests against the +immutable [35 GB Zenodo archive](https://zenodo.org/api/records/18738598/files/lms.zip/content) +to retrieve only the compressed English member: **8,582,666,912 bytes** +(approximately 8.0 GiB). It streams the raw DEFLATE data into a temporary file, +producing a **14,600,342,241-byte** model (approximately 13.6 GiB), and checks +the member metadata, CRC-32, expanded-model SHA-256 +(`d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00`), +and KenLM readability before atomically installing it. The server or any proxy +must support standards-compliant byte ranges (HTTP 206 and `Content-Range`); +the helper refuses an HTTP 200 response rather than accidentally downloading +the complete archive. Allow space for the installed model, its temporary +expanded copy, and at least a 1 GiB safety margin. The record is openly +accessible and licensed **CC BY 4.0**, which permits automatic download and +redistribution with attribution. Without the option, a missing model still +causes DArtP to be reported as skipped. ArtP does not need the model. Custom +standalone file/archive mirrors remain supported but must be supplied together with +`--language-model-sha256`; the project-specific environment equivalents are +`PATHBENCH_LANGUAGE_MODEL_URL`, `PATHBENCH_LANGUAGE_MODEL_SHA256`, and +`PATHBENCH_LANGUAGE_MODEL_CACHE`. The built-in range mode retains no compressed +archive and treats the verified model under `lms/` as its cache; the cache +directory applies to custom downloads. + +The built-in Zenodo artifact identity used by the live integration is: + +| Field | Published value | +| --- | --- | +| Archive (`lms.zip`) size | `35017940434` bytes | +| Member | `lms/wiki_en_token.arpa.bin` | +| Decompressed member size | `14600342241` bytes | +| Compressed member size | `8582666912` bytes | +| Member CRC-32 | `5afb90ef` | +| Decompressed member SHA-256 | `d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00` | + +The SHA-256 value is specifically the digest of the fully transferred and +decompressed `lms/wiki_en_token.arpa.bin` member. It is **not** a digest of +`lms.zip` or of the member's compressed DEFLATE stream. + +For both tests together, allow **at least 12 GB of system RAM and 8 GB of GPU +VRAM**; **16 GB system RAM and 12–16 GB VRAM are recommended** to leave room +for both wav2vec2 models, the decoder, and transient activations. Any NVIDIA +CUDA GPU supported by the selected PyTorch wheel is acceptable; a T4 (16 GB), +L4 (24 GB), A10/A10G (24 GB), V100 (16/32 GB), or A100 works. Smaller 8 GB +cards may require closing other GPU processes and can run out of memory on +long audio. AMD ROCm GPUs, Apple GPUs, and CPU-only runtimes do not satisfy +this CUDA smoke test. + +Google Colab GPUs can be used. Select a GPU runtime and confirm that +`nvidia-smi` works; the commonly assigned T4 and higher-memory L4/A100 options +meet the recommendation. Colab does not guarantee a particular GPU, RAM +amount, availability, or uninterrupted runtime, and its temporary filesystem +means the environment and downloaded checkpoints may need to be recreated in +a later session. If Colab assigns a smaller GPU or low-RAM runtime, inspect +`nvidia-smi` and available system memory before running the helper. +On a Python 3.12 T4 runtime, use PyTorch 2.6.0's CUDA 12.4 wheels: + +```bash +python tools/test_gpu_predictors.py --download-language-model --cuda-version 12.4 +``` + +A successful run ends with `2 passed`. A second invocation reuses both +`tools/gpu_venv` and the verified model cache rather than downloading them +again. Colab's local disk is ephemeral, however, so the installed 13.6 GiB +model and environment are lost when its runtime is recycled. Persist the +checkout itself if reuse across sessions is important (`PATHBENCH_LANGUAGE_MODEL_CACHE` +only controls custom standalone downloads). + **Without sudo access:** A containerised environment such as Docker is recommended. ## Downloads @@ -213,7 +443,18 @@ find /path/to/your/datasets/easycall/EasyCall/m13 -name "m13 _*" -exec bash -c ' ### N-gram models -The n-gram models required for DArtP and ArtP are included in the [Oral Cancer - YouTube](https://zenodo.org/records/18738598) download. +The n-gram models required by DArtP are included in the +[Oral Cancer - YouTube](https://zenodo.org/records/18738598) download. ArtP does +not require an n-gram model. Create `lms/` at the repository root and copy the +models there with these exact names: + +| Language | Filename | +| --- | --- | +| English | `wiki_en_token.arpa` or `wiki_en_token.arpa.bin` | +| Dutch | `wiki_nl_token.arpa` or `wiki_nl_token.arpa.bin` | +| Spanish | `wiki_es_token.arpa.bin` | +| Italian | `wiki_it_token.arpa.bin` | +| Mandarin Chinese | `wiki_zh_token.arpa` or `wiki_zh_token.arpa.bin` | ## Testing @@ -228,6 +469,19 @@ python -m pytest tests/test_evaluators.py::TestEvaluatorMethods -v All tests should pass. If all evaluator tests fail simultaneously, the reference audio file in `tests/data/test_audio.wav` may be corrupted — the `test_audio_integrity` test will confirm this. +To test only ArtP and DArtP after downloading the English n-gram model, run: + +```bash +python -m pytest \ + tests/test_evaluators.py::TestEvaluatorMethods::test_articulatory_precision \ + tests/test_evaluators.py::TestEvaluatorMethods::test_artp_double_asr -v +``` + +The first run downloads the Hugging Face checkpoints used by the phonetic and +English ASR models and therefore requires network access and several gigabytes +of free disk space. The DArtP test is reported as skipped, rather than failed, +when neither English n-gram filename listed above exists. + > **Note:** During the NAD evaluator tests you will see a `Wav2Vec2Model LOAD REPORT` table listing several keys (e.g. `project_q`, `quantizer`) as **UNEXPECTED**. These warnings are harmless — the keys belong to pre-training heads that are not needed for feature extraction and can be safely ignored. ### Dataset integrity @@ -280,4 +534,3 @@ This work is partly financed by the Dutch Research Council (NWO) under project n ## Author Bence Mark Halpern, Nagoya University - diff --git a/docs/installation.rst b/docs/installation.rst index c855f59..b3bfe30 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -39,5 +39,8 @@ Without sudo access, a containerised environment such as Docker is recommended. .. note:: - PathBench cannot be published to PyPI because it depends on Git-hosted forks - of ``phonemizer`` and ``pyctcdecode``. + PathBench's Python dependencies are available as versioned package-index + releases, including ``phonemizer-fork==3.3.2`` and + ``pyctcdecode==0.5.0``. Installing them does not require GitHub credentials. + The ``espeak-ng`` shared library remains a separate system dependency; use + the revision documented in the project README for reproducible IPA output. diff --git a/pyproject.toml b/pyproject.toml index 641655f..83ec385 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,8 +50,10 @@ dependencies = [ "transformers", "dtw-python", "jiwer", - "phonemizer-fork @ git+https://github.com/thewh1teagle/phonemizer-fork.git", - "pyctcdecode @ git+https://github.com/kensho-technologies/pyctcdecode.git", + # PyPI release of the fork; it continues to expose the ``phonemizer`` API. + "phonemizer-fork==3.3.2", + # Keep every runtime dependency installable without VCS/GitHub access. + "pyctcdecode==0.5.0", "praat-parselmouth>=0.4.4", "scikit-learn", ] diff --git a/tests/test_gpu_predictors_tool.py b/tests/test_gpu_predictors_tool.py new file mode 100644 index 0000000..4b3a5a9 --- /dev/null +++ b/tests/test_gpu_predictors_tool.py @@ -0,0 +1,572 @@ +"""Network-free tests for the opt-in GPU smoke-test model downloader.""" + +from __future__ import annotations + +from contextlib import nullcontext +import hashlib +import importlib.util +import io +import os +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace +import urllib.error +import zipfile +import zlib + +import pytest + + +SPEC = importlib.util.spec_from_file_location( + "test_gpu_predictors", Path(__file__).parents[1] / "tools/test_gpu_predictors.py" +) +assert SPEC and SPEC.loader +tool = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = tool +SPEC.loader.exec_module(tool) + + +def test_builtin_language_model_checksum_matches_published_value(): + assert tool.LANGUAGE_MODEL_SHA256 == ( + "d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00" + ) + + +class Response(io.BytesIO): + status = 206 + + def __init__(self, value=b"", headers=None, status=206): + super().__init__(value) + self.headers = headers or {} + self.status = status + + def getcode(self): + return self.status + + def __enter__(self): + return self + + def __exit__(self, *_args): + self.close() + + +def zip_bytes(files: dict[str, bytes]) -> bytes: + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as archive: + for name, contents in files.items(): + archive.writestr(name, contents) + return output.getvalue() + + +def configure(monkeypatch, tmp_path: Path, payload: bytes): + calls = [] + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path / "repo") + monkeypatch.setattr( + tool.urllib.request, "urlopen", + lambda *_args, **_kwargs: calls.append(True) or Response(payload), + ) + return calls, hashlib.sha256(payload).hexdigest() + + +def test_successful_download_prefers_binary_and_reuses_cache(monkeypatch, tmp_path): + payload = zip_bytes({"models/wiki_en_token.arpa": b"text", "wiki_en_token.arpa.bin": b"binary"}) + calls, digest = configure(monkeypatch, tmp_path, payload) + first = tool.install_language_model("https://example/model.zip", digest, tmp_path / "cache") + second = tool.install_language_model("https://example/model.zip", digest, tmp_path / "cache") + assert first.read_bytes() == second.read_bytes() == b"binary" + assert len(calls) == 1 + + +def test_installing_arpa_removes_stale_preferred_binary(monkeypatch, tmp_path): + payload = zip_bytes({"models/wiki_en_token.arpa": b"replacement arpa"}) + _calls, digest = configure(monkeypatch, tmp_path, payload) + stale_binary = tmp_path / "repo" / "lms" / "wiki_en_token.arpa.bin" + stale_binary.parent.mkdir(parents=True) + stale_binary.write_bytes(b"stale binary") + + installed = tool.install_language_model( + "https://example/model.zip", digest, tmp_path / "cache", + ) + + assert installed.name == "wiki_en_token.arpa" + assert installed.read_bytes() == b"replacement arpa" + assert not stale_binary.exists() + + +def test_checksum_mismatch_removes_partial(monkeypatch, tmp_path): + calls, _ = configure(monkeypatch, tmp_path, b"not expected") + with pytest.raises(RuntimeError, match="SHA-256 mismatch"): + tool.install_language_model("https://example/model", "0" * 64, tmp_path / "cache") + assert calls and not list((tmp_path / "cache").iterdir()) + + +def test_interrupted_download_removes_partial(monkeypatch, tmp_path): + class Interrupted(Response): + def read(self, *_args): + raise OSError("connection reset") + + monkeypatch.setattr(tool.urllib.request, "urlopen", lambda *_a, **_k: Interrupted(b"x")) + with pytest.raises(RuntimeError, match="download failed"): + tool.install_language_model("https://example/model", "0" * 64, tmp_path / "cache") + assert not list((tmp_path / "cache").iterdir()) + + +def test_http_failure(monkeypatch, tmp_path): + def fail(*_args, **_kwargs): + raise urllib.error.HTTPError("url", 503, "unavailable", {}, None) + + monkeypatch.setattr(tool.urllib.request, "urlopen", fail) + with pytest.raises(RuntimeError, match="503"): + tool.install_language_model("https://example/model", "0" * 64, tmp_path / "cache") + + +def test_safe_archive_extracts_only_model(tmp_path): + archive = tmp_path / "models.zip" + archive.write_bytes(zip_bytes({"docs/readme": b"no", "nested/wiki_en_token.arpa": b"yes"})) + output = tmp_path / "out" + output.mkdir() + assert tool._extract_model(archive, output).read_bytes() == b"yes" + assert sorted(path.name for path in output.iterdir()) == ["wiki_en_token.arpa"] + + +@pytest.mark.parametrize("name", ["../wiki_en_token.arpa", "/tmp/wiki_en_token.arpa", "C:\\..\\wiki_en_token.arpa"]) +def test_malicious_archive_path_is_rejected(tmp_path, name): + archive = tmp_path / "bad.zip" + archive.write_bytes(zip_bytes({name: b"bad"})) + output = tmp_path / "out" + output.mkdir() + with pytest.raises(RuntimeError, match="unsafe"): + tool._extract_model(archive, output) + + +def test_corrupt_kenlm_data_is_reported(monkeypatch, tmp_path): + model = tmp_path / "wiki_en_token.arpa.bin" + model.write_bytes(b"corrupt") + monkeypatch.setattr(tool.subprocess, "run", lambda *_a, **_k: (_ for _ in ()).throw( + tool.subprocess.CalledProcessError(1, "python") + )) + with pytest.raises(tool.CommandError, match="Validating the English KenLM model failed"): + tool.validate_language_model(Path("python"), model) + + +def test_already_installed_verified_model(monkeypatch, tmp_path): + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path) + model = tmp_path / "lms" / "wiki_en_token.arpa.bin" + model.parent.mkdir() + model.write_bytes(b"model") + assert tool.installed_language_model(hashlib.sha256(b"model").hexdigest()) == model + assert tool.installed_language_model("0" * 64) is None + + +def test_rejected_preferred_binary_blocks_arpa_fallback(monkeypatch, tmp_path): + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path) + models = tmp_path / "lms" + models.mkdir() + (models / "wiki_en_token.arpa.bin").write_bytes(b"rejected binary") + (models / "wiki_en_token.arpa").write_bytes(b"otherwise valid arpa") + + assert tool.installed_language_model("0" * 64) is None + + +def test_rejected_preferred_binary_is_replaced(monkeypatch, tmp_path): + rejected = tmp_path / "lms" / "wiki_en_token.arpa.bin" + fallback = tmp_path / "lms" / "wiki_en_token.arpa" + replacement = tmp_path / "replacement" / "wiki_en_token.arpa.bin" + rejected.parent.mkdir() + rejected.write_bytes(b"rejected binary") + fallback.write_bytes(b"otherwise valid arpa") + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path) + monkeypatch.setattr( + tool, + "install_zenodo_language_model", + lambda **_kwargs: replacement, + ) + + assert tool.prepare_language_model( + download=True, + url=tool.LANGUAGE_MODEL_URL, + expected_sha256="0" * 64, + cache_dir=tmp_path / "cache", + force=False, + installed_sha256="0" * 64, + ) == replacement + + +def test_forced_download_replaces_an_installed_model(monkeypatch, tmp_path): + installed = tmp_path / "lms" / "wiki_en_token.arpa.bin" + refreshed = tmp_path / "refreshed" / "wiki_en_token.arpa.bin" + monkeypatch.setattr(tool, "installed_language_model", lambda _digest: installed) + calls = [] + + def install(url, digest, cache, *, force): + calls.append((url, digest, cache, force)) + return refreshed + + monkeypatch.setattr(tool, "install_language_model", install) + result = tool.prepare_language_model( + download=True, + url="https://example/custom-model", + expected_sha256="1" * 64, + cache_dir=tmp_path / "cache", + force=True, + installed_sha256=None, + ) + + assert result == refreshed + assert calls == [ + ("https://example/custom-model", "1" * 64, tmp_path / "cache", True) + ] + + +def test_force_without_download_remains_opted_out(monkeypatch, tmp_path): + installed = tmp_path / "lms" / "wiki_en_token.arpa.bin" + monkeypatch.setattr(tool, "installed_language_model", lambda _digest: installed) + monkeypatch.setattr( + tool, + "install_language_model", + lambda *_args, **_kwargs: pytest.fail("download must remain opt-in"), + ) + + assert tool.prepare_language_model( + download=False, + url="https://example/custom-model", + expected_sha256="1" * 64, + cache_dir=tmp_path / "cache", + force=True, + installed_sha256=None, + ) == installed + + +def test_manual_binary_is_not_compared_with_download_digest(): + assert tool.installed_model_digest( + download=False, + url=tool.LANGUAGE_MODEL_URL, + expected_sha256=tool.LANGUAGE_MODEL_SHA256, + ) is None + + +def test_managed_builtin_download_verifies_installed_binary(): + assert tool.installed_model_digest( + download=True, + url=tool.LANGUAGE_MODEL_URL, + expected_sha256=tool.LANGUAGE_MODEL_SHA256, + ) == tool.LANGUAGE_MODEL_SHA256 + + +def test_download_remains_opt_in(monkeypatch): + monkeypatch.setattr(sys, "argv", ["tool"]) + assert tool.parse_args().download_language_model is False + + +def test_native_install_remains_opt_in_and_cli_overrides_environment(monkeypatch): + monkeypatch.setenv("PATHBENCH_INSTALL_SYSTEM_DEPENDENCIES", "1") + monkeypatch.setattr(sys, "argv", ["tool", "--no-install-system-dependencies"]) + assert tool.parse_args().install_system_dependencies is False + + +class FakeNative: + def __init__(self, *, release=None, uid=0, programs=None, responses=None): + self.release = release or {"ID": "ubuntu"} + self.uid = uid + self.programs = {"apt-get": "/usr/bin/apt-get", **(programs or {})} + self.responses = list(responses or []) + self.commands = [] + + def os_release(self): return self.release + def which(self, name): return self.programs.get(name) + def is_colab(self): return False + def is_container(self): return False + def is_virtual_machine(self): return False + def execute(self, command, **_kwargs): + self.commands.append(command) + if self.responses: + return self.responses.pop(0) + return tool.subprocess.CompletedProcess(command, 0, "", "") + def privilege_prefix(self): + if self.uid == 0: + return [] + sudo = self.which("sudo") + if sudo and self.execute([sudo, "-n", "true"]).returncode == 0: + return [sudo, "-n"] + raise RuntimeError("no changes were made") + + +def test_supported_ubuntu_and_unsupported_distribution(): + assert tool.supported_apt_host(FakeNative()) + assert not tool.supported_apt_host(FakeNative(release={"ID": "fedora"})) + + +def test_root_and_passwordless_sudo_prefixes(): + assert FakeNative(uid=0).privilege_prefix() == [] + fake = FakeNative(uid=1000, programs={"sudo": "/usr/bin/sudo"}) + assert fake.privilege_prefix() == ["/usr/bin/sudo", "-n"] + assert fake.commands == [["/usr/bin/sudo", "-n", "true"]] + + +def test_missing_escalation_fails_before_changes(): + fake = FakeNative(uid=1000) + with pytest.raises(RuntimeError, match="no changes"): + fake.privilege_prefix() + assert fake.commands == [] + + +def test_one_apt_install_contains_expected_packages(monkeypatch): + fake = FakeNative() + monkeypatch.setattr(tool.Path, "resolve", lambda self: Path("/opt/uv/python")) + tool.install_system_packages(fake, "python3.12") + installs = [c for c in fake.commands if "install" in c] + assert len(installs) == 1 + assert set(tool.SYSTEM_PACKAGES) <= set(installs[0]) + assert "python3-venv" not in installs[0] + + +def test_failed_venv_creation_installs_python3_venv(monkeypatch): + fake = FakeNative(responses=[subprocess.CompletedProcess([], 1)]) + monkeypatch.setattr(tool.Path, "resolve", lambda self: Path("/usr/bin/python3")) + + tool.install_system_packages(fake, "python3") + + assert fake.commands[0][:3] == ["python3", "-m", "venv"] + installs = [command for command in fake.commands if "install" in command] + assert "python3-venv" in installs[0] + + +def test_apt_failure_preserves_status(monkeypatch): + fake = FakeNative(responses=[tool.subprocess.CompletedProcess([], 7)]) + monkeypatch.setattr(tool.Path, "resolve", lambda self: Path("/opt/uv/python")) + with pytest.raises(tool.CommandError) as error: + tool.install_system_packages(fake, "python") + assert error.value.returncode == 7 + + +def test_matching_espeak_marker_is_reused(monkeypatch, tmp_path): + marker = tmp_path / "marker" + marker.write_text(tool.ESPEAK_NG_COMMIT) + monkeypatch.setattr(tool, "espeak_probe", lambda _system: True) + fake = FakeNative() + tool.ensure_espeak(fake, install=False, marker=marker) + assert fake.commands == [] + + +def test_mismatched_marker_requires_opt_in(tmp_path): + marker = tmp_path / "marker" + marker.write_text("old") + with pytest.raises(RuntimeError, match=tool.ESPEAK_NG_COMMIT): + tool.ensure_espeak(FakeNative(), install=False, marker=marker) + + +def test_espeak_install_does_not_fetch_abbreviated_commit(monkeypatch, tmp_path): + fake = FakeNative() + monkeypatch.setattr( + tool.tempfile, "TemporaryDirectory", lambda **_kwargs: nullcontext(str(tmp_path)) + ) + monkeypatch.setattr(tool, "espeak_probe", lambda _system: True) + tool.ensure_espeak(fake, install=True, marker=tmp_path / "marker") + assert not any(command[:4] == ["git", "-C", str(tmp_path / "source"), "fetch"] + for command in fake.commands) + assert ["git", "-C", str(tmp_path / "source"), "checkout", "--detach", + tool.ESPEAK_NG_COMMIT] in fake.commands + + +@pytest.mark.parametrize( + "colab,container,virtual_machine", + [(True, False, False), (False, True, False), (False, False, True)], +) +def test_driver_install_refused_in_managed_environments( + monkeypatch, colab, container, virtual_machine, +): + fake = FakeNative() + monkeypatch.setattr(fake, "is_colab", lambda: colab) + monkeypatch.setattr(fake, "is_container", lambda: container) + monkeypatch.setattr(fake, "is_virtual_machine", lambda: virtual_machine) + with pytest.raises(RuntimeError, match="host"): + tool.nvidia_driver_setup(fake, confirm=True) + assert fake.commands == [] + + +def test_driver_header_query_uses_valid_status_format(): + fake = FakeNative(responses=[ + subprocess.CompletedProcess([], 0, "", ""), + subprocess.CompletedProcess([], 0, "6.8.0-test\n", ""), + subprocess.CompletedProcess([], 0, "install ok installed", ""), + subprocess.CompletedProcess([], 0, "driver : nvidia-driver-550 - distro non-free recommended\n", ""), + ]) + tool.nvidia_driver_setup(fake, confirm=False) + assert fake.commands[2] == ["dpkg-query", "-W", "-f=${Status}", + "linux-headers-6.8.0-test"] + + +def test_driver_only_main_skips_unrelated_espeak_validation(monkeypatch, tmp_path): + args = SimpleNamespace(venv=tmp_path / "venv", install_nvidia_driver=True, + confirm_nvidia_driver_install=False) + monkeypatch.setattr(tool, "parse_args", lambda: args) + monkeypatch.setattr(tool, "nvidia_driver_setup", lambda *_args, **_kwargs: False) + monkeypatch.setattr(tool, "ensure_espeak", lambda *_args, **_kwargs: + pytest.fail("driver-only setup must not validate espeak")) + assert tool.main(FakeNative()) == 0 + + +@pytest.mark.parametrize("name", ["APPTAINER_NAME", "SINGULARITY_NAME", "container"]) +def test_native_system_detects_additional_container_environments(monkeypatch, name): + monkeypatch.setenv(name, "managed") + assert tool.NativeSystem(which=lambda _name: None).is_container() + + +def test_native_system_uses_vm_specific_virtualization_probe(): + commands = [] + + def execute(command, **_kwargs): + commands.append(command) + return subprocess.CompletedProcess(command, 0) + + system = tool.NativeSystem(runner=execute, which=lambda name: f"/usr/bin/{name}") + assert system.is_virtual_machine() + assert commands == [["/usr/bin/systemd-detect-virt", "--quiet", "--vm"]] + + +def test_range_reader_validates_content_range_and_coalesces(monkeypatch): + payload = bytes(range(256)) * 10000 + calls = [] + + def open_range(request, **_kwargs): + start, end = map(int, request.headers["Range"].removeprefix("bytes=").split("-")) + calls.append((start, end)) + return Response(payload[start:end + 1], { + "Content-Range": f"bytes {start}-{end}/{len(payload)}" + }) + + monkeypatch.setattr(tool.urllib.request, "urlopen", open_range) + reader = tool.BufferedHTTPRangeReader( + tool.HTTPRangeClient("https://example/archive", len(payload)), 1024 * 1024 + ) + reader.seek(17) + assert reader.read(5) == payload[17:22] + reader.seek(800_000) + assert reader.read(5) == payload[800_000:800_005] + assert len(calls) == 1 + + +@pytest.mark.parametrize("header", [None, "bytes 0-8/10", "nonsense"]) +def test_malformed_content_range_is_rejected(monkeypatch, header): + monkeypatch.setattr(tool.urllib.request, "urlopen", lambda *_a, **_k: Response( + b"0123456789", {"Content-Range": header} if header else {} + )) + with pytest.raises(RuntimeError, match="Content-Range"): + tool.HTTPRangeClient("https://example/archive", 10).read(0, 9) + + +def test_http_200_range_response_is_rejected(monkeypatch): + monkeypatch.setattr(tool.urllib.request, "urlopen", lambda *_a, **_k: Response( + b"whole archive", status=200 + )) + with pytest.raises(RuntimeError, match="HTTP 200"): + tool.HTTPRangeClient("https://example/archive", 100).read(0, 9) + + +def _compressed_case(data: bytes): + compressor = zlib.compressobj(wbits=-zlib.MAX_WBITS) + compressed = compressor.compress(data) + compressor.flush() + + class Client: + def chunks(self, _start, _end): + yield compressed + + return Client(), compressed + + +def test_raw_deflate_member_download(tmp_path): + data = b"ordinary ZIP member" * 100 + client, compressed = _compressed_case(data) + output = tmp_path / "model" + tool._expand_deflate_range( + client, 0, len(compressed) - 1, output, + compressed_size=len(compressed), expanded_size=len(data), + expected_crc=zlib.crc32(data), expected_sha256=hashlib.sha256(data).hexdigest(), + ) + assert output.read_bytes() == data + + +@pytest.mark.parametrize("failure", ["truncated", "overlong", "crc", "sha"]) +def test_corrupt_member_downloads_are_rejected(tmp_path, failure): + data = b"model data" * 100 + client, compressed = _compressed_case(data) + kwargs = dict(compressed_size=len(compressed), expanded_size=len(data), + expected_crc=zlib.crc32(data), + expected_sha256=hashlib.sha256(data).hexdigest()) + if failure == "truncated": + kwargs["compressed_size"] += 1 + elif failure == "overlong": + kwargs["expanded_size"] -= 1 + elif failure == "crc": + kwargs["expected_crc"] ^= 1 + else: + kwargs["expected_sha256"] = "0" * 64 + with pytest.raises(RuntimeError): + tool._expand_deflate_range(client, 0, len(compressed) - 1, + tmp_path / failure, **kwargs) + + +def test_zip64_extra_supports_large_sizes(): + large_expanded, large_compressed = 6_000_000_000, 5_000_000_000 + extra = b"\x01\x00\x10\x00" + large_expanded.to_bytes(8, "little") \ + + large_compressed.to_bytes(8, "little") + assert tool._zip64_values(extra) == [large_expanded, large_compressed] + + +def test_zip_metadata_derives_member_bounds_and_checks_local_header(monkeypatch): + data = b"member" * 100 + output = io.BytesIO() + with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr(tool.LANGUAGE_MODEL_MEMBER, data) + payload = bytearray(output.getvalue()) + + class MemoryClient: + size = len(payload) + + def read(self, start, end): + return bytes(payload[start:end + 1]) + + with zipfile.ZipFile(io.BytesIO(payload)) as archive: + expected = archive.getinfo(tool.LANGUAGE_MODEL_MEMBER) + monkeypatch.setattr(tool, "LANGUAGE_MODEL_SIZE", len(data)) + monkeypatch.setattr(tool, "LANGUAGE_MODEL_COMPRESSED_SIZE", expected.compress_size) + monkeypatch.setattr(tool, "LANGUAGE_MODEL_CRC32", expected.CRC) + info, start, end = tool.inspect_zenodo_member(MemoryClient()) + assert end - start + 1 == info.compress_size + + payload[info.header_offset + 30] ^= 1 + with pytest.raises(RuntimeError, match="local and central"): + tool.inspect_zenodo_member(MemoryClient()) + + +def test_insufficient_disk_space(monkeypatch, tmp_path): + usage = type("Usage", (), {"free": 1})() + monkeypatch.setattr(tool.shutil, "disk_usage", lambda _path: usage) + with pytest.raises(RuntimeError, match="Insufficient disk space"): + tool._check_model_space(tmp_path) + + +def test_retry_exhaustion(monkeypatch): + calls = [] + monkeypatch.setattr(tool.time, "sleep", lambda _seconds: None) + + def fail(*_args, **_kwargs): + calls.append(1) + raise urllib.error.URLError("interrupted") + + monkeypatch.setattr(tool.urllib.request, "urlopen", fail) + with pytest.raises(RuntimeError, match="range request failed"): + tool.HTTPRangeClient("https://example/archive", 10, retries=2).read(0, 9) + assert len(calls) == 3 + + +@pytest.mark.skipif(not os.environ.get("PATHBENCH_LIVE_ZENODO_METADATA"), + reason="set PATHBENCH_LIVE_ZENODO_METADATA=1 for live range check") +def test_live_zenodo_metadata(): + client = tool.HTTPRangeClient(tool.LANGUAGE_MODEL_URL, + tool.LANGUAGE_MODEL_ARCHIVE_SIZE) + info, start, end = tool.inspect_zenodo_member(client) + assert info.filename == tool.LANGUAGE_MODEL_MEMBER + assert end - start + 1 == tool.LANGUAGE_MODEL_COMPRESSED_SIZE + assert (start, end) == (18_177_078_384, 26_759_745_295) diff --git a/tests/test_string_clean.py b/tests/test_string_clean.py new file mode 100644 index 0000000..002801e --- /dev/null +++ b/tests/test_string_clean.py @@ -0,0 +1,37 @@ +"""Tests for text cleaning and the pinned phonemizer API.""" + +import importlib.util +import os +from pathlib import Path + +import pytest + + +_SPEC = importlib.util.spec_from_file_location( + "pathbench_string_clean", Path(__file__).parents[1] / "pathbench" / "string_clean.py" +) +assert _SPEC and _SPEC.loader +_STRING_CLEAN = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_STRING_CLEAN) +cached_phonemize = _STRING_CLEAN.cached_phonemize + + +def test_phonemizer_fork_exposes_expected_api(): + """The PyPI fork must retain the imports used by ``string_clean``.""" + from phonemizer.phonemize import phonemize + from phonemizer.separator import Separator + + assert callable(phonemize) + assert Separator(phone=" ", word="|").phone == " " + + +@pytest.mark.skipif( + os.environ.get("PATHBENCH_TEST_PINNED_ESPEAK") != "1", + reason="requires the README-pinned espeak-ng commit", +) +def test_pinned_espeak_preserves_language_specific_ipa(): + """The documented espeak-ng revision uses an Italian tap, unlike English.""" + cached_phonemize.cache_clear() + + assert cached_phonemize("Roma", "it").strip() == "ɾ o m a" + assert cached_phonemize("Roma", "en-us").strip() == "ɹ oʊ m ə" diff --git a/tools/test_gpu_predictors.py b/tools/test_gpu_predictors.py new file mode 100755 index 0000000..6aef554 --- /dev/null +++ b/tools/test_gpu_predictors.py @@ -0,0 +1,988 @@ +#!/usr/bin/env python3 +"""Install a CUDA-enabled PathBench environment and smoke-test ArtP and DArtP. + +Stepwise installation (each step skips components that are already suitable): + +1. Install the Ubuntu prerequisites listed in README.md if they are missing. +2. Build the pinned espeak-ng commit unless the installation marker confirms it. +3. Ensure an NVIDIA driver is installed and ``nvidia-smi --list-gpus`` works. + Host driver installation is intentionally left to the machine or cloud provider. +4. From the PathBench checkout, run ``python3 tools/test_gpu_predictors.py``. +5. This script then reuses or creates ``tools/gpu_venv``; installs PyTorch, + torchaudio, PathBench, and test dependencies only when its import/version + checks fail; verifies CUDA access; and runs the focused ArtP and DArtP tests. +6. Download the English n-gram model described in README.md to exercise DArtP; + without it, pytest reports the DArtP test as skipped. + +Run with ``--help`` to select another Python, CUDA/PyTorch version, or venv. +Allow at least 12 GB system RAM and 8 GB GPU VRAM (16 GB RAM and 12--16 GB +VRAM recommended). Google Colab NVIDIA GPU runtimes are supported when +``nvidia-smi`` works; T4, L4, and A100 assignments meet the recommendation. +""" + +from __future__ import annotations + +import argparse +import ctypes.util +import hashlib +import http.client +import io +import os +from pathlib import Path +import shutil +import subprocess +import sys +import tarfile +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +import zipfile +import binascii +import re +import struct +import zlib + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent +MODEL_NAMES = ("wiki_en_token.arpa.bin", "wiki_en_token.arpa") +# The versioned record URL is deliberately not the mutable ``latest`` link. +LANGUAGE_MODEL_URL = ( + "https://zenodo.org/api/records/18738598/files/lms.zip/content" +) +# The MD5 describes the complete archive. Range extraction cannot verify it. +LANGUAGE_MODEL_ARCHIVE_SIZE = 35_017_940_434 +LANGUAGE_MODEL_ARCHIVE_MD5 = "01d62027902e93270e5f0d00806c473c" +LANGUAGE_MODEL_MEMBER = "lms/wiki_en_token.arpa.bin" +LANGUAGE_MODEL_SIZE = 14_600_342_241 +LANGUAGE_MODEL_COMPRESSED_SIZE = 8_582_666_912 +LANGUAGE_MODEL_CRC32 = 0x5AFB90EF +# Calculated from the decompressed member, not lms.zip or its DEFLATE stream. +LANGUAGE_MODEL_SHA256 = "d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00" +USER_AGENT = "PathBench GPU predictor model installer/1.0" +RANGE_BLOCK_SIZE = 1024 * 1024 +RANGE_RETRIES = 4 +ESPEAK_NG_COMMIT = "2ea41210" +ESPEAK_NG_MARKER = Path("/usr/local/share/pathbench/espeak-ng-commit") +SYSTEM_PACKAGES = ( + "git", "ca-certificates", "curl", "build-essential", "cmake", + "ninja-build", "pkg-config", "libfftw3-dev", "liblapack-dev", +) +MANUAL_PREREQUISITES = ", ".join(SYSTEM_PACKAGES) + ", and espeak-ng commit " + ESPEAK_NG_COMMIT + + +class NativeSystem: + """Small injectable boundary around native-host operations.""" + + def __init__(self, *, runner=None, which=None, geteuid=None) -> None: + self.runner = runner or subprocess.run + self.which = which or shutil.which + self.geteuid = geteuid or getattr(os, "geteuid", lambda: 1) + + def execute(self, command, **kwargs): + return self.runner(command, **kwargs) + + def os_release(self) -> dict[str, str]: + values = {} + try: + for line in Path("/etc/os-release").read_text().splitlines(): + if "=" in line: + key, value = line.split("=", 1) + values[key] = value.strip().strip('"') + except OSError: + pass + return values + + def is_colab(self) -> bool: + return bool(os.environ.get("COLAB_RELEASE_TAG") or Path("/content").exists()) + + def is_container(self) -> bool: + marker_paths = ( + Path("/.dockerenv"), + Path("/run/.containerenv"), + Path("/run/systemd/container"), + Path("/.singularity.d"), + ) + container_environment = ( + "container", + "KUBERNETES_SERVICE_HOST", + "APPTAINER_NAME", + "SINGULARITY_NAME", + ) + if any(path.exists() for path in marker_paths) or any( + os.environ.get(name) for name in container_environment + ): + return True + try: + cgroups = Path("/proc/1/cgroup").read_text() + except OSError: + cgroups = "" + if re.search(r"/(?:docker|lxc|libpod|podman|kubepods)(?:[-/.]|$)", cgroups): + return True + detector = self.which("systemd-detect-virt") + return bool(detector and self.execute( + [detector, "--quiet", "--container"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ).returncode == 0) + + def is_virtual_machine(self) -> bool: + """Return whether systemd identifies this host as a virtual machine.""" + detector = self.which("systemd-detect-virt") + return bool(detector and self.execute( + [detector, "--quiet", "--vm"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ).returncode == 0) + + def privilege_prefix(self) -> list[str]: + if self.geteuid() == 0: + return [] + sudo = self.which("sudo") + if sudo and self.execute([sudo, "-n", "true"], stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL).returncode == 0: + return [sudo, "-n"] + raise RuntimeError("System installation requires root or passwordless noninteractive sudo; no changes were made.") + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} + + +class CommandError(RuntimeError): + """A child command failed and its exit status should be preserved.""" + + def __init__(self, message: str, returncode: int) -> None: + super().__init__(message) + self.returncode = returncode + + +def run( + command: list[str], description: str, *, cwd: Path | None = None +) -> None: + """Run a visible command and turn a nonzero status into a useful error.""" + print(f"\n==> {description}", flush=True) + print("+ " + " ".join(command), flush=True) + try: + subprocess.run(command, check=True, cwd=cwd) + except subprocess.CalledProcessError as error: + raise CommandError( + f"{description} failed with exit status {error.returncode}. " + f"Review the command output above: {' '.join(command)}", + error.returncode, + ) from error + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Create a CUDA-enabled PathBench virtual environment, verify GPU " + "access, and run the ArtP and DArtP smoke tests." + ) + ) + parser.add_argument( + "--python", default=os.environ.get("PYTHON", "python3"), + help="Python 3.10-3.12 executable (default: PYTHON or python3)", + ) + parser.add_argument( + "--pytorch-version", default=os.environ.get("PYTORCH_VERSION", "2.6.0"), + help="matching PyTorch and torchaudio version (default: 2.6.0)", + ) + parser.add_argument( + "--cuda-version", default=os.environ.get("PATHBENCH_CUDA_VERSION", "12.4"), + help="CUDA wheel version, such as 12.4 (default: 12.4)", + ) + parser.add_argument( + "--venv", type=Path, + default=Path(os.environ.get("VENV", SCRIPT_DIR / "gpu_venv")), + help="virtual-environment destination (default: tools/gpu_venv)", + ) + parser.add_argument( + "--download-language-model", action="store_true", + help="download and install the English DArtP model (large; opt in)", + ) + parser.add_argument( + "--language-model-url", + default=os.environ.get("PATHBENCH_LANGUAGE_MODEL_URL", LANGUAGE_MODEL_URL), + help="model/archive URL (PATHBENCH_LANGUAGE_MODEL_URL)", + ) + parser.add_argument( + "--language-model-sha256", + default=os.environ.get("PATHBENCH_LANGUAGE_MODEL_SHA256"), + help="SHA-256 for a custom URL (PATHBENCH_LANGUAGE_MODEL_SHA256)", + ) + parser.add_argument( + "--language-model-cache", type=Path, + default=Path(os.environ.get("PATHBENCH_LANGUAGE_MODEL_CACHE", "~/.cache/pathbench")), + help="download cache (PATHBENCH_LANGUAGE_MODEL_CACHE; default: ~/.cache/pathbench)", + ) + parser.add_argument( + "--force-language-model-download", action="store_true", + help="discard a cached artifact and download it again", + ) + parser.add_argument( + "--install-system-dependencies", action=argparse.BooleanOptionalAction, + default=_env_flag("PATHBENCH_INSTALL_SYSTEM_DEPENDENCIES"), + help="install native Debian/Ubuntu prerequisites (PATHBENCH_INSTALL_SYSTEM_DEPENDENCIES)", + ) + parser.add_argument( + "--install-nvidia-driver", action=argparse.BooleanOptionalAction, + default=_env_flag("PATHBENCH_INSTALL_NVIDIA_DRIVER"), + help="EXPERIMENTAL: preflight a driver install on administrator-controlled bare-metal Ubuntu", + ) + parser.add_argument( + "--confirm-nvidia-driver-install", action="store_true", + help="confirm the separately requested Ubuntu-recommended driver installation", + ) + return parser.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +class HTTPRangeClient: + """Strict, retrying HTTP byte-range client for a pinned immutable object.""" + + def __init__(self, url: str, size: int, *, retries: int = RANGE_RETRIES) -> None: + self.url = url + self.size = size + self.retries = retries + + def _open(self, start: int, end: int): + if start < 0 or end < start or end >= self.size: + raise RuntimeError(f"Invalid archive byte range {start}-{end}") + request = urllib.request.Request( + self.url, + headers={"Range": f"bytes={start}-{end}", "User-Agent": USER_AGENT}, + ) + response = urllib.request.urlopen(request, timeout=60) + status = getattr(response, "status", None) or response.getcode() + if status != 206: + response.close() + raise RuntimeError( + f"Range server returned HTTP {status}, not 206; refusing a possible " + f"full {self.size:,}-byte archive response" + ) + value = response.headers.get("Content-Range") + match = re.fullmatch(r"bytes (\d+)-(\d+)/(\d+)", value or "") + if not match: + response.close() + raise RuntimeError(f"Malformed or missing Content-Range: {value!r}") + reported = tuple(map(int, match.groups())) + if reported != (start, end, self.size): + response.close() + raise RuntimeError( + f"Unexpected Content-Range {value!r}; expected bytes " + f"{start}-{end}/{self.size}" + ) + return response + + @staticmethod + def _transient(error: BaseException) -> bool: + if isinstance(error, RuntimeError): + return False + if isinstance(error, urllib.error.HTTPError): + return error.code in {408, 425, 429, 500, 502, 503, 504} + return isinstance(error, (OSError, http.client.HTTPException, + urllib.error.URLError)) + + def read(self, start: int, end: int) -> bytes: + for attempt in range(self.retries + 1): + try: + with self._open(start, end) as response: + data = response.read(end - start + 1) + if response.read(1) or len(data) != end - start + 1: + raise OSError("truncated or overlong HTTP range response") + return data + except Exception as error: + if not self._transient(error) or attempt == self.retries: + raise RuntimeError(f"HTTP range request failed: {error}") from error + time.sleep(min(2 ** attempt, 8)) + raise AssertionError("unreachable") + + def chunks(self, start: int, end: int, chunk_size: int = 4 * 1024 * 1024): + """Yield an exact interval, resuming interrupted responses with a new range.""" + position = start + failures = 0 + while position <= end: + response = None + try: + response = self._open(position, end) + while position <= end: + chunk = response.read(min(chunk_size, end - position + 1)) + if not chunk: + raise OSError("truncated HTTP range response") + position += len(chunk) + yield chunk + if response.read(1): + raise OSError("overlong HTTP range response") + failures = 0 + except Exception as error: + if not self._transient(error) or failures >= self.retries: + raise RuntimeError(f"HTTP range download failed: {error}") from error + time.sleep(min(2 ** failures, 8)) + failures += 1 + finally: + if response is not None: + response.close() + + +class BufferedHTTPRangeReader(io.RawIOBase): + """Seekable range-backed file with block coalescing for ZIP metadata reads.""" + + def __init__(self, client: HTTPRangeClient, block_size: int = RANGE_BLOCK_SIZE): + self.client = client + self.block_size = block_size + self.position = 0 + self.cache: dict[int, bytes] = {} + + def readable(self): + return True + + def seekable(self): + return True + + def tell(self): + return self.position + + def seek(self, offset, whence=os.SEEK_SET): + positions = {os.SEEK_SET: offset, os.SEEK_CUR: self.position + offset, + os.SEEK_END: self.client.size + offset} + if whence not in positions or positions[whence] < 0: + raise ValueError("invalid seek") + self.position = positions[whence] + return self.position + + def read(self, size=-1): + if size is None or size < 0: + size = self.client.size - self.position + size = min(size, self.client.size - self.position) + output = bytearray() + while size > 0: + block = self.position // self.block_size + if block not in self.cache: + start = block * self.block_size + end = min(start + self.block_size, self.client.size) - 1 + self.cache[block] = self.client.read(start, end) + data = self.cache[block] + within = self.position - block * self.block_size + take = min(size, len(data) - within) + output.extend(data[within:within + take]) + self.position += take + size -= take + return bytes(output) + + +def _zip64_values(extra: bytes) -> list[int]: + position = 0 + while position + 4 <= len(extra): + kind, length = struct.unpack_from(" tuple[zipfile.ZipInfo, int, int]: + """Read ZIP/ZIP64 metadata remotely and return the raw DEFLATE interval.""" + reader = BufferedHTTPRangeReader(client) + try: + with zipfile.ZipFile(reader) as archive: + info = archive.getinfo(LANGUAGE_MODEL_MEMBER) + except (KeyError, zipfile.BadZipFile, OSError) as error: + raise RuntimeError(f"Cannot inspect Zenodo ZIP metadata: {error}") from error + if not _safe_member(info.filename) or info.filename != LANGUAGE_MODEL_MEMBER: + raise RuntimeError(f"Unsafe or unexpected ZIP member name: {info.filename!r}") + expected = (zipfile.ZIP_DEFLATED, LANGUAGE_MODEL_SIZE, + LANGUAGE_MODEL_COMPRESSED_SIZE, LANGUAGE_MODEL_CRC32) + actual = (info.compress_type, info.file_size, info.compress_size, info.CRC) + if actual != expected: + raise RuntimeError(f"Zenodo member metadata changed: {actual!r} != {expected!r}") + if info.flag_bits & 1: + raise RuntimeError("Encrypted ZIP members are unsupported") + + fixed = client.read(info.header_offset, info.header_offset + 29) + signature, _version, flags, method, _time, _date, crc, compressed, expanded, name_len, extra_len = \ + struct.unpack(" bool: + path = Path(name.replace("\\", "/")) + return not path.is_absolute() and ".." not in path.parts + + +def _extract_model(archive: Path, destination: Path) -> Path: + """Extract only the preferred English model, after checking every path.""" + members: dict[str, object] + opener: object + if zipfile.is_zipfile(archive): + opener = zipfile.ZipFile(archive) + members = {item.filename: item for item in opener.infolist()} + elif tarfile.is_tarfile(archive): + opener = tarfile.open(archive) + members = {item.name: item for item in opener.getmembers()} + else: + # A direct ARPA or KenLM binary needs no extraction. + target = destination / next( + (name for name in MODEL_NAMES if name in archive.name), MODEL_NAMES[0] + ) + shutil.copyfile(archive, target) + return target + with opener: + if any(not _safe_member(name) for name in members): + raise RuntimeError("Language-model archive contains an unsafe path") + selected = next( + (name for wanted in MODEL_NAMES for name in members if Path(name).name == wanted), + None, + ) + if selected is None: + raise RuntimeError("Archive contains neither wiki_en_token.arpa.bin nor wiki_en_token.arpa") + target = destination / Path(selected).name + source = opener.open(members[selected]) if isinstance(opener, zipfile.ZipFile) else opener.extractfile(members[selected]) + if source is None: + raise RuntimeError("Selected language-model archive member is not a file") + with source, target.open("wb") as output: + shutil.copyfileobj(source, output) + return target + + +def _check_model_space(directory: Path, required_size: int = LANGUAGE_MODEL_SIZE) -> None: + """Require room for the staged model plus a conservative one-GiB margin.""" + available = shutil.disk_usage(directory).free + required = required_size + 1024 ** 3 + if available < required: + raise RuntimeError( + f"Insufficient disk space: {available:,} bytes free; {required:,} required " + "for the temporary expanded model and safety margin" + ) + + +def _expand_deflate_range( + client: HTTPRangeClient, start: int, end: int, destination: Path, *, + compressed_size: int = LANGUAGE_MODEL_COMPRESSED_SIZE, + expanded_size: int = LANGUAGE_MODEL_SIZE, + expected_crc: int = LANGUAGE_MODEL_CRC32, + expected_sha256: str = LANGUAGE_MODEL_SHA256, +) -> None: + """Download and authenticate one raw-DEFLATE ZIP member.""" + inflater = zlib.decompressobj(-zlib.MAX_WBITS) + transferred = expanded = crc = 0 + digest = hashlib.sha256() + last_report = time.monotonic() + with destination.open("wb") as output: + for chunk in client.chunks(start, end): + transferred += len(chunk) + if transferred > compressed_size: + raise RuntimeError("Compressed member exceeds its advertised size") + data = inflater.decompress(chunk, expanded_size - expanded + 1) + if inflater.unconsumed_tail: + raise RuntimeError("Expanded model exceeds its advertised size") + expanded += len(data) + if expanded > expanded_size: + raise RuntimeError("Expanded model exceeds its advertised size") + output.write(data) + crc = binascii.crc32(data, crc) + digest.update(data) + if time.monotonic() - last_report >= 5: + print(f"Transferred {transferred:,}/{compressed_size:,}; expanded " + f"{expanded:,}/{expanded_size:,} bytes...", flush=True) + last_report = time.monotonic() + tail = inflater.flush() + expanded += len(tail) + if expanded > expanded_size: + raise RuntimeError("Expanded model exceeds its advertised size") + output.write(tail) + crc = binascii.crc32(tail, crc) + digest.update(tail) + print(f"Transferred {transferred:,}/{compressed_size:,}; expanded " + f"{expanded:,}/{expanded_size:,} bytes.", flush=True) + if transferred != compressed_size: + raise RuntimeError(f"Compressed size mismatch: {transferred:,} != {compressed_size:,}") + if not inflater.eof or inflater.unused_data: + raise RuntimeError("Truncated or trailing compressed member data") + if expanded != expanded_size: + raise RuntimeError(f"Expanded size mismatch: {expanded:,} != {expanded_size:,}") + if crc & 0xFFFFFFFF != expected_crc: + raise RuntimeError(f"Model CRC-32 mismatch: {crc & 0xFFFFFFFF:08x}") + actual_sha256 = digest.hexdigest() + if actual_sha256 != expected_sha256.lower(): + raise RuntimeError( + f"Model SHA-256 mismatch (expected {expected_sha256}, got {actual_sha256})" + ) + + +def install_zenodo_language_model( + *, force: bool = False, validator=None, +) -> Path: + """Range-extract, validate, and atomically install the pinned ZIP member.""" + models_dir = REPO_ROOT / "lms" + models_dir.mkdir(parents=True, exist_ok=True) + installed = models_dir / Path(LANGUAGE_MODEL_MEMBER).name + if not force and installed.is_file() and installed.stat().st_size == LANGUAGE_MODEL_SIZE \ + and sha256(installed) == LANGUAGE_MODEL_SHA256: + print(f"Reusing verified installed language model: {installed}") + return installed + print( + "The Zenodo range download will transfer approximately 8.0 GiB and install " + "a 13.6 GiB English language model. Additional temporary space and a 1 GiB " + "safety margin are required.", flush=True, + ) + _check_model_space(models_dir) + client = HTTPRangeClient(LANGUAGE_MODEL_URL, LANGUAGE_MODEL_ARCHIVE_SIZE) + _info, start, end = inspect_zenodo_member(client) + temporary = Path(tempfile.mkstemp(prefix=".wiki-en-", dir=models_dir)[1]) + try: + _expand_deflate_range(client, start, end, temporary) + if validator is not None: + validator(temporary) + os.replace(temporary, installed) + return installed + finally: + temporary.unlink(missing_ok=True) + + +def install_language_model( + url: str, expected_sha256: str, cache_dir: Path, *, force: bool = False, + validator=None, +) -> Path: + """Install a custom standalone file/archive with a user-supplied digest.""" + if not expected_sha256 or len(expected_sha256) != 64: + raise RuntimeError("A pinned 64-character SHA-256 is required for the language model") + expected_sha256 = expected_sha256.lower() + cache_dir = cache_dir.expanduser().resolve() + cache_dir.mkdir(parents=True, exist_ok=True) + source_name = Path(urllib.parse.unquote(urllib.parse.urlparse(url).path)).name + source_name = source_name if source_name else "download" + artifact = cache_dir / f"language-model-{expected_sha256}-{source_name}" + if force: + artifact.unlink(missing_ok=True) + if artifact.exists() and sha256(artifact) != expected_sha256: + artifact.unlink() + if not artifact.exists(): + temporary = Path(tempfile.mkstemp(prefix=".language-model-", dir=cache_dir)[1]) + try: + request = urllib.request.Request(url, headers={"User-Agent": "PathBench/0.1"}) + with urllib.request.urlopen(request, timeout=60) as response, temporary.open("wb") as output: + total = 0 + last_report = time.monotonic() + while chunk := response.read(1024 * 1024): + output.write(chunk) + total += len(chunk) + if time.monotonic() - last_report >= 5: + print(f"Downloaded {total:,} bytes...", flush=True) + last_report = time.monotonic() + print(f"Downloaded {total:,} bytes.", flush=True) + actual = sha256(temporary) + if actual != expected_sha256: + raise RuntimeError( + f"Language-model SHA-256 mismatch (expected {expected_sha256}, got {actual})" + ) + os.replace(temporary, artifact) + except (OSError, urllib.error.URLError, http.client.HTTPException) as error: + raise RuntimeError(f"Language-model download failed: {error}") from error + finally: + temporary.unlink(missing_ok=True) + else: + print(f"Reusing verified language-model download: {artifact}") + + models_dir = REPO_ROOT / "lms" + models_dir.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".model-", dir=models_dir)) + try: + extracted = _extract_model(artifact, staging) + if extracted.stat().st_size == 0: + raise RuntimeError("Downloaded language model is empty") + if validator is not None: + validator(extracted) + installed = models_dir / extracted.name + if installed.name == "wiki_en_token.arpa": + # The evaluator prefers the binary whenever it exists. Removing a + # stale binary ensures it actually consumes the verified ARPA we + # are about to install and validate. + (models_dir / "wiki_en_token.arpa.bin").unlink(missing_ok=True) + os.replace(extracted, installed) + return installed + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def installed_language_model(expected_sha256: str | None = None) -> Path | None: + """Return the preferred installed model, rejecting empty/known-bad files.""" + for name in MODEL_NAMES: + path = REPO_ROOT / "lms" / name + if not path.is_file() or path.stat().st_size == 0: + continue + # The built-in artifact is the compiled model itself, so its digest also + # authenticates an already-installed copy. Archive digests do not. + if expected_sha256 and name == "wiki_en_token.arpa.bin": + if sha256(path) != expected_sha256.lower(): + # The evaluator always prefers this binary when it exists. Do + # not fall back to an ARPA that the smoke test would not use; + # make the opted-in preparation path replace the bad binary. + return None + return path + return None + + +def validate_language_model(venv_python: Path, model: Path) -> None: + """Have the target environment parse the model before expensive tests start.""" + run([ + str(venv_python), "-c", + "import kenlm, pathlib, sys; p=pathlib.Path(sys.argv[1]); " + "assert p.stat().st_size, 'model is empty'; kenlm.Model(str(p))", + str(model), + ], "Validating the English KenLM model") + + +def prepare_language_model( + *, + download: bool, + url: str, + expected_sha256: str, + cache_dir: Path, + force: bool, + installed_sha256: str | None, + validator=None, +) -> Path | None: + """Reuse an installed model, unless an opted-in forced refresh was requested.""" + model = installed_language_model(installed_sha256) + if download and (model is None or force): + if url == LANGUAGE_MODEL_URL: + return install_zenodo_language_model(force=force, validator=validator) + kwargs = {"force": force} + if validator is not None: + kwargs["validator"] = validator + return install_language_model(url, expected_sha256, cache_dir, **kwargs) + return model + + +def installed_model_digest( + *, download: bool, url: str, expected_sha256: str | None, +) -> str | None: + """Return a digest only when this run is managing the pinned artifact. + + Manually installed models are supported and need not be byte-for-byte + identical to the project's default Zenodo binary. They are checked by + KenLM's parser instead of against the download artifact's digest. + """ + if download and url == LANGUAGE_MODEL_URL: + return expected_sha256 + return None + + +def require_program(program: str, explanation: str) -> str: + path = shutil.which(program) + if path is None: + raise RuntimeError(f"Required program '{program}' was not found. {explanation}") + return path + + +def python_succeeds(python: Path | str, code: str) -> bool: + """Return whether an environment already satisfies a Python-side check.""" + return subprocess.run( + [str(python), "-c", code], stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode == 0 + + +def supported_apt_host(system: NativeSystem) -> bool: + release = system.os_release() + identities = {release.get("ID", ""), *release.get("ID_LIKE", "").split()} + return bool(identities & {"ubuntu", "debian"} and system.which("apt-get")) + + +def _checked(system: NativeSystem, command: list[str], description: str, **kwargs): + print(f"\n==> {description}\n+ {' '.join(map(str, command))}", flush=True) + result = system.execute(command, **kwargs) + if result.returncode: + raise CommandError(f"{description} failed with exit status {result.returncode}: " + + " ".join(map(str, command)), result.returncode) + return result + + +def install_system_packages(system: NativeSystem, python: str) -> list[str]: + """Install build prerequisites only after the caller's explicit opt-in.""" + if not supported_apt_host(system): + raise RuntimeError("Automatic native setup supports only Ubuntu/Debian with apt-get. " + f"Install these prerequisites manually: {MANUAL_PREREQUISITES}") + prefix = system.privilege_prefix() # establish access before apt update changes state + packages = list(SYSTEM_PACKAGES) + distro_python = Path(python).resolve().as_posix().startswith("/usr/bin/python") + if distro_python: + # Importing venv does not prove that Debian's separately packaged + # ensurepip payload is installed. Exercise the operation we need on a + # disposable directory so fresh/minimal hosts get python3-venv. + with tempfile.TemporaryDirectory(prefix="pathbench-venv-probe-") as probe: + if system.execute( + [python, "-m", "venv", str(Path(probe) / "venv")], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ).returncode: + packages.append("python3-venv") + env = os.environ.copy() + env["DEBIAN_FRONTEND"] = "noninteractive" + print("Native packages to install: " + ", ".join(packages)) + _checked(system, prefix + ["apt-get", "update"], "Updating apt package metadata", env=env) + _checked(system, prefix + ["apt-get", "install", "-y", "--no-install-recommends", *packages], + "Installing native prerequisites", env=env) + return prefix + + +def espeak_probe(system: NativeSystem) -> bool: + executable = system.which("espeak-ng") + if not executable or not ctypes.util.find_library("espeak-ng"): + return False + return system.execute([executable, "--ipa", "-q", "PathBench"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL).returncode == 0 + + +def ensure_espeak(system: NativeSystem, *, install: bool, + privilege_prefix: list[str] | None = None, + marker: Path = ESPEAK_NG_MARKER) -> None: + marker_value = None + try: + marker_value = marker.read_text().strip() + except OSError: + pass + if marker_value == ESPEAK_NG_COMMIT and espeak_probe(system): + print(f"\n==> Reusing pinned espeak-ng {ESPEAK_NG_COMMIT}") + return + if not install: + raise RuntimeError( + f"Pinned espeak-ng revision {ESPEAK_NG_COMMIT} is required; marker {marker} " + "is missing/mismatched or its runtime probe failed. Re-run with " + "--install-system-dependencies to rebuild it." + ) + prefix = privilege_prefix if privilege_prefix is not None else system.privilege_prefix() + jobs = max(1, min(os.cpu_count() or 1, 8)) + with tempfile.TemporaryDirectory(prefix="pathbench-espeak-") as temporary: + source, build = Path(temporary) / "source", Path(temporary) / "build" + _checked(system, ["git", "clone", "https://github.com/espeak-ng/espeak-ng.git", str(source)], + "Cloning espeak-ng") + _checked(system, ["git", "-C", str(source), "checkout", "--detach", ESPEAK_NG_COMMIT], + "Checking out pinned espeak-ng revision") + _checked(system, ["cmake", "-S", str(source), "-B", str(build), "-G", "Ninja", + "-DUSE_ASYNC=OFF", "-DBUILD_SHARED_LIBS=ON"], "Configuring espeak-ng") + _checked(system, ["cmake", "--build", str(build), "--parallel", str(jobs)], "Building espeak-ng") + _checked(system, prefix + ["cmake", "--install", str(build)], "Installing espeak-ng") + _checked(system, prefix + ["ldconfig"], "Refreshing the shared-library cache") + if not espeak_probe(system): + raise RuntimeError("The installed espeak-ng failed its executable/shared-library phonemization probe; marker was not written.") + staged = Path(temporary) / "espeak-ng-commit" + staged.write_text(ESPEAK_NG_COMMIT + "\n") + _checked(system, prefix + ["install", "-D", "-m", "0644", str(staged), str(marker) + ".tmp"], + "Staging the espeak-ng revision marker") + _checked(system, prefix + ["mv", str(marker) + ".tmp", str(marker)], + "Recording the espeak-ng revision atomically") + + +def nvidia_driver_setup(system: NativeSystem, *, confirm: bool) -> bool: + """Preflight/install Ubuntu's recommended driver; return reboot-required.""" + if (system.is_colab() or system.is_container() or system.is_virtual_machine() + or os.environ.get("WSL_DISTRO_NAME")): + raise RuntimeError( + "NVIDIA driver installation is refused in virtual machines, Colab, WSL, " + "containers, Kubernetes, and GPU-passthrough environments. Install the driver " + "on the host; restarting a " + "notebook/container cannot activate a newly installed host kernel module." + ) + if system.os_release().get("ID") != "ubuntu": + raise RuntimeError("NVIDIA driver installation is supported only on a bare-metal Ubuntu host.") + _checked(system, ["lspci"], "Identifying NVIDIA hardware") + kernel = system.execute(["uname", "-r"], text=True, capture_output=True) + if kernel.returncode: + raise CommandError("Reading the running kernel failed", kernel.returncode) + kernel_version = kernel.stdout.strip() + headers = f"linux-headers-{kernel_version}" + header_check = system.execute(["dpkg-query", "-W", "-f=${Status}", headers], + text=True, capture_output=True) + if header_check.returncode or "ok installed" not in header_check.stdout: + raise RuntimeError(f"Matching running-kernel headers are required: {headers}") + devices = system.execute(["ubuntu-drivers", "devices"], text=True, capture_output=True) + if devices.returncode: + raise CommandError("Querying Ubuntu's recommended NVIDIA driver failed", devices.returncode) + matches = re.findall(r"(nvidia-driver-\d+(?:-server)?)\s+.*recommended", devices.stdout) + if not matches: + raise RuntimeError("ubuntu-drivers did not report a recommended signed NVIDIA driver package.") + package = matches[0] + print(f"Running kernel: {kernel_version}\nUbuntu recommended driver: {package}") + if not confirm: + print("Dry run only; add --confirm-nvidia-driver-install to install this package.") + return False + prefix = system.privilege_prefix() + env = os.environ.copy() + env["DEBIAN_FRONTEND"] = "noninteractive" + _checked(system, prefix + ["apt-get", "update"], "Updating apt package metadata", env=env) + _checked(system, prefix + ["apt-get", "install", "-y", package], + "Installing Ubuntu-recommended signed NVIDIA driver", env=env) + print("NVIDIA driver installed successfully. A host reboot is required; stopping now.") + return True + + +def main(system: NativeSystem | None = None) -> int: + args = parse_args() + system = system or NativeSystem() + venv = args.venv.expanduser().resolve() + try: + if args.install_nvidia_driver: + nvidia_driver_setup(system, confirm=args.confirm_nvidia_driver_install) + return 0 + python = require_program( + args.python, "Set --python to a Python 3.10-3.12 executable." + ) + privilege_prefix = None + if args.install_system_dependencies: + privilege_prefix = install_system_packages(system, python) + ensure_espeak(system, install=args.install_system_dependencies, + privilege_prefix=privilege_prefix) + nvidia_smi = require_program( + "nvidia-smi", "Install an NVIDIA driver and expose the GPU to this environment." + ) + + version_check = subprocess.run( + [ + python, "-c", + "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}'); " + "raise SystemExit(0 if (3, 10) <= sys.version_info[:2] <= (3, 12) else 1)", + ], + text=True, capture_output=True, + ) + if version_check.returncode != 0: + detected = version_check.stdout.strip() or "unknown" + raise RuntimeError( + f"Python 3.10-3.12 is required, but '{python}' reports {detected}." + ) + + run([nvidia_smi, "--list-gpus"], "Checking the NVIDIA driver and visible GPUs") + if not (venv / "bin" / "python").is_file(): + run([python, "-m", "venv", str(venv)], "Creating the GPU virtual environment") + else: + print(f"\n==> Reusing existing virtual environment: {venv}") + venv_python = venv / "bin" / "python" + cuda_tag = f"cu{args.cuda_version.replace('.', '')}" + + torch_check = ( + "import torch, torchaudio; " + f"assert torch.__version__.split('+')[0] == '{args.pytorch_version}'; " + f"assert torchaudio.__version__.split('+')[0] == '{args.pytorch_version}'; " + f"assert (torch.version.cuda or '').replace('.', '') == '{cuda_tag[2:]}'" + ) + if not python_succeeds(venv_python, torch_check): + run([ + str(venv_python), "-m", "pip", "install", + "--force-reinstall", + f"torch=={args.pytorch_version}", f"torchaudio=={args.pytorch_version}", + "--index-url", f"https://download.pytorch.org/whl/{cuda_tag}", + ], "Installing missing or mismatched CUDA-enabled PyTorch packages") + else: + print("\n==> Matching CUDA-enabled PyTorch packages are already installed") + + project_check = ( + "from pathlib import Path; import pathbench, pytest, kenlm; " + f"assert Path(pathbench.__file__).resolve().is_relative_to(Path({str(REPO_ROOT)!r}))" + ) + if not python_succeeds(venv_python, project_check): + run([ + str(venv_python), "-m", "pip", "install", "-e", f"{REPO_ROOT}[all]", + "pytest", "kenlm", + ], "Installing missing PathBench or test dependencies") + else: + print("\n==> PathBench and test dependencies are already installed") + run([ + str(venv_python), "-c", + "import torch; " + "print(f'PyTorch: {torch.__version__}'); " + "print(f'PyTorch CUDA runtime: {torch.version.cuda}'); " + "assert torch.cuda.is_available(), " + "'CUDA build installed, but no GPU is available; check the driver and container access'; " + "print(f'GPU: {torch.cuda.get_device_name(0)}')", + ], "Verifying that PyTorch can use the GPU") + + expected_digest = args.language_model_sha256 + if args.language_model_url == LANGUAGE_MODEL_URL and not expected_digest: + expected_digest = LANGUAGE_MODEL_SHA256 + known_installed_digest = installed_model_digest( + download=args.download_language_model, + url=args.language_model_url, + expected_sha256=expected_digest, + ) + model = prepare_language_model( + download=args.download_language_model, + url=args.language_model_url, + expected_sha256=expected_digest or "", + cache_dir=args.language_model_cache, + force=args.force_language_model_download, + installed_sha256=known_installed_digest, + validator=lambda path: validate_language_model(venv_python, path), + ) + if model is None: + print( + "\nWarning: the English n-gram model is absent. The DArtP test " + "will be reported as skipped; see README.md's N-gram models section.", + file=sys.stderr, + ) + if model is not None: + validate_language_model(venv_python, model) + + pytest_environment = os.environ.copy() + pytest_environment["MPLBACKEND"] = "Agg" + print("\n==> Running the ArtP and DArtP smoke tests", flush=True) + command = [ + str(venv_python), "-m", "pytest", + "tests/test_evaluators.py::TestEvaluatorMethods::test_articulatory_precision", + "tests/test_evaluators.py::TestEvaluatorMethods::test_artp_double_asr", "-v", + ] + print("+ " + " ".join(command), flush=True) + try: + subprocess.run(command, check=True, cwd=REPO_ROOT, env=pytest_environment) + except subprocess.CalledProcessError as error: + raise CommandError( + f"Running the ArtP and DArtP smoke tests failed with exit status " + f"{error.returncode}. Review the command output above: {' '.join(command)}", + error.returncode, + ) from error + except CommandError as error: + print(f"\nError: {error}", file=sys.stderr) + return error.returncode + except RuntimeError as error: + print(f"\nError: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())