Files
imagefind/scripts/live-speech-acceptance.py

300 lines
12 KiB
Python

#!/usr/bin/env python3
"""Verify live audio indexing with a run-owned video containing clear speech."""
from __future__ import annotations
import argparse
import atexit
import hashlib
import json
import time
from pathlib import Path
from typing import Any
import httpx
def call(client: httpx.Client, method: str, path: str, **kwargs: Any) -> httpx.Response:
response = client.request(method, path, **kwargs)
if response.is_error:
raise RuntimeError(f"{method} {path} -> {response.status_code}: {response.text[:1000]}")
return response
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
while block := handle.read(1024 * 1024):
digest.update(block)
return digest.hexdigest()
def find_video(client: httpx.Client, source_id: str, filename: str) -> dict[str, Any] | None:
videos = call(client, "GET", "/api/v1/videos?limit=500").json()
return next(
(
item
for item in videos
if item.get("source_id") == source_id
and item.get("source_key", "").rsplit("/", 1)[-1] == filename
),
None,
)
def upload_fixture(
client: httpx.Client,
fixture: Path,
state: dict[str, Any],
) -> tuple[dict[str, Any], list[float]]:
filename = f"{state['run_id']}-{fixture.name}"
uploads = call(client, "GET", "/api/v1/uploads?limit=500").json()
upload = next((item for item in uploads if item.get("filename") == filename), None)
if upload is None or upload.get("status") in {"failed", "cancelled"}:
upload = call(
client,
"POST",
"/api/v1/uploads",
json={
"source_id": state["source_id"],
"relative_path": "ingest",
"filename": filename,
"title": "音频 GPU 验收",
"collection_id": state["collection_id"],
"tag_ids": [state["tag_id"]],
"size_bytes": fixture.stat().st_size,
"sha256": sha256(fixture),
"conflict": "skip",
},
).json()
latencies: list[float] = []
if upload.get("status") != "completed":
hidden = find_video(client, state["source_id"], filename)
if hidden is not None:
raise AssertionError("unfinished speech upload leaked into /videos")
chunk_size = int(upload["chunk_size"])
received = {int(index) for index in upload.get("received_chunks", upload.get("received", []))}
with fixture.open("rb") as handle:
for index in range(int(upload["total_chunks"])):
data = handle.read(chunk_size)
if index in received:
continue
started = time.perf_counter()
call(
client,
"PUT",
f"/api/v1/uploads/{upload['id']}/chunks/{index}",
content=data,
headers={"X-Chunk-SHA256": hashlib.sha256(data).hexdigest()},
)
latencies.append(round((time.perf_counter() - started) * 1000, 1))
upload_id = upload["id"]
upload = call(client, "POST", f"/api/v1/uploads/{upload_id}/complete").json()
else:
upload_id = upload["id"]
deadline = time.monotonic() + 240
while upload.get("status") not in {"completed", "failed", "cancelled"}:
if time.monotonic() >= deadline:
raise TimeoutError(f"speech upload did not finish: {upload.get('status')}")
time.sleep(1)
uploads = call(client, "GET", "/api/v1/uploads?limit=500").json()
upload = next(item for item in uploads if item["id"] == upload_id)
if upload["status"] != "completed":
raise AssertionError(f"speech upload failed: {upload.get('message')}")
return upload, latencies
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--base-url", required=True)
auth = parser.add_mutually_exclusive_group(required=True)
auth.add_argument("--token-file", type=Path)
auth.add_argument("--password")
parser.add_argument("--run-dir", type=Path, required=True)
parser.add_argument("--fixture", default="e2e-speech.mp4")
parser.add_argument(
"--query",
action="append",
dest="queries",
help="Required audio-search keyword; repeat to verify multiple words (default: Americans)",
)
parser.add_argument("--timeout", type=int, default=1200)
args = parser.parse_args()
state_path = args.run_dir / "state.json"
state = json.loads(state_path.read_text())
fixture = args.run_dir / args.fixture
if not fixture.is_file():
raise FileNotFoundError(fixture)
headers = {}
if args.token_file is not None:
headers["Authorization"] = f"Bearer {args.token_file.read_text().strip()}"
client = httpx.Client(
base_url=args.base_url.rstrip("/"),
headers=headers,
timeout=httpx.Timeout(60, connect=10),
)
if args.password is not None:
login = call(
client,
"POST",
"/api/v1/auth/login",
json={"password": args.password, "remember_device": False},
).json()
csrf = str(login.get("csrf_token") or "")
if not csrf:
raise AssertionError("login response omitted CSRF token")
client.headers["X-CSRF-Token"] = csrf
def logout() -> None:
try:
client.post("/api/v1/auth/logout")
except Exception:
pass
atexit.register(logout)
# The repository's default live fixture contains the JFK inauguration line
# "And so, my fellow Americans...". Keep the default assertion tied to
# the actual fixture; callers using another recording can repeat --query.
required_queries = [value.strip() for value in (args.queries or ["Americans"]) if value.strip()]
if not required_queries:
raise ValueError("at least one non-empty --query is required")
initial_models = call(client, "GET", "/api/v1/models").json()
print(
json.dumps(
{
"audio_model": (initial_models.get("manifest") or {})
.get("components", {})
.get("audio"),
"audio_health": (initial_models.get("component_health") or {}).get("audio"),
"audio_accelerator": (initial_models.get("accelerator") or {})
.get("components", {})
.get("audio"),
},
ensure_ascii=False,
),
flush=True,
)
upload, chunk_latencies = upload_fixture(client, fixture, state)
filename = upload["filename"]
deadline = time.monotonic() + args.timeout
observed_devices: list[str] = []
latest_message = "等待视频入库"
positives: dict[str, dict[str, Any]] = {}
video: dict[str, Any] | None = None
last_report_at = 0.0
while time.monotonic() < deadline:
video = find_video(client, state["source_id"], filename)
models = call(client, "GET", "/api/v1/models").json()
audio = models.get("accelerator", {}).get("components", {}).get("audio", {})
device = audio.get("actual_device") or audio.get("device")
if device and device not in observed_devices:
observed_devices.append(device)
if video is not None:
for query in required_queries:
if query in positives:
continue
result = call(
client,
"POST",
"/api/v1/search",
json={"text": query, "recognition_types": ["audio"]},
).json()
positive = next(
(item for item in result.get("items", []) if item.get("video_id") == video["id"]),
None,
)
if positive is not None:
positives[query] = positive
if len(positives) == len(required_queries):
break
jobs = call(client, "GET", "/api/v1/jobs?page=1&page_size=10").json()["items"]
audio_jobs = [item for item in jobs if item.get("kind") == "transcribe_audio"]
if audio_jobs:
latest_message = audio_jobs[0].get("message") or audio_jobs[0].get("status", "")
if audio_jobs[0].get("status") == "failed":
raise AssertionError(f"audio job failed: {audio_jobs[0].get('error')}")
now = time.monotonic()
if now - last_report_at >= 10:
print(
json.dumps(
{
"waiting": latest_message,
"video_found": video is not None,
"observed_devices": observed_devices,
},
ensure_ascii=False,
),
flush=True,
)
last_report_at = now
time.sleep(2)
if video is None:
raise AssertionError("speech video never appeared in /videos")
missing_queries = [query for query in required_queries if query not in positives]
if missing_queries:
raise TimeoutError(
"audio search never matched required keywords "
f"{missing_queries}; matched={sorted(positives)}; latest={latest_message}"
)
for query, positive in positives.items():
if positive.get("segment_start_ms") is None or positive.get("segment_end_ms") is None:
raise AssertionError(f"audio match omitted time segment for {query!r}: {positive}")
match_details = positive.get("match_details", [])
audio_text = " ".join(
str(detail.get("text") or "")
for detail in match_details
if detail.get("type") == "audio"
).strip()
if not audio_text:
raise AssertionError(f"audio match omitted recognized text for {query!r}: {positive}")
if len([character for character in audio_text if character.isalnum()]) < 8:
raise AssertionError(
f"audio transcript is implausibly short for {query!r}: {audio_text!r}"
)
negative = call(
client,
"POST",
"/api/v1/search",
json={"text": "pineapple submarine 9834721", "recognition_types": ["audio"]},
).json()
if any(item.get("video_id") == video["id"] for item in negative.get("items", [])):
raise AssertionError("non-matching audio search unexpectedly returned the speech video")
jobs = call(client, "GET", "/api/v1/jobs?page=1&page_size=10").json()["items"]
latest_audio = next((item for item in jobs if item.get("kind") == "transcribe_audio"), None)
report = {
"upload_id": upload["id"],
"video_id": video["id"],
"chunk_latency_ms": chunk_latencies,
"observed_audio_devices": observed_devices,
"audio_job": latest_audio,
"matches": {
query: {
"segment_start_ms": positive.get("segment_start_ms"),
"segment_end_ms": positive.get("segment_end_ms"),
"details": positive.get("match_details", []),
}
for query, positive in positives.items()
},
"negative_matches": len(negative.get("items", [])),
}
if not any(str(device).upper().startswith("GPU") for device in observed_devices):
raise AssertionError(f"GPU was not observed during speech indexing: {observed_devices}")
state[f"speech_upload_id_{fixture.stem}"] = upload["id"]
state[f"speech_video_id_{fixture.stem}"] = video["id"]
state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n")
(args.run_dir / "speech-report.json").write_text(
json.dumps(report, ensure_ascii=False, indent=2) + "\n"
)
print(json.dumps(report, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()