Files
imagefind/scripts/live-final-acceptance.py
T

295 lines
12 KiB
Python

#!/usr/bin/env python3
"""Cover final live gaps without touching records outside the supplied run prefix."""
from __future__ import annotations
import argparse
import hashlib
import json
import math
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
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 percentile(values: list[float], ratio: float) -> float:
ordered = sorted(values)
return ordered[max(0, math.ceil(len(ordered) * ratio) - 1)] if ordered else 0.0
def find_video(client: httpx.Client, source_id: str, filename: str) -> dict[str, Any] | None:
return next(
(
video
for video in call(client, "GET", "/api/v1/videos?limit=500").json()
if video.get("source_id") == source_id
and str(video.get("source_key") or "").rsplit("/", 1)[-1] == filename
),
None,
)
def wait_upload(client: httpx.Client, upload_id: str, timeout: int = 180) -> dict[str, Any]:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
upload = next(
(item for item in call(client, "GET", "/api/v1/uploads?limit=500").json() if item["id"] == upload_id),
None,
)
if upload and upload["status"] in {"completed", "failed", "cancelled"}:
return upload
time.sleep(1)
raise TimeoutError(f"upload did not reach a terminal state: {upload_id}")
def wait_job(client: httpx.Client, job_id: str, timeout: int = 180) -> dict[str, Any]:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
job = next(
(item for item in call(client, "GET", "/api/v1/jobs?limit=500").json() if item["id"] == job_id),
None,
)
if job and job["status"] in {"completed", "failed", "cancelled"}:
return job
time.sleep(1)
raise TimeoutError(f"job did not reach a terminal state: {job_id}")
def performance_sample(base_url: str, client: httpx.Client) -> dict[str, Any]:
paths = (
"/api/v1/status",
"/api/v1/videos?limit=50",
"/api/v1/jobs?page=1&page_size=10",
"/api/v1/profile",
"/api/v1/models",
"/api/v1/system/resources",
)
def request(path: str) -> tuple[str, float, int]:
started = time.perf_counter()
response = httpx.get(
base_url + path,
headers=dict(client.headers),
cookies=dict(client.cookies),
timeout=httpx.Timeout(20, connect=5),
)
return path, round((time.perf_counter() - started) * 1000, 1), response.status_code
samples: dict[str, list[float]] = {path: [] for path in paths}
errors: list[dict[str, Any]] = []
with ThreadPoolExecutor(max_workers=4) as pool:
futures = [pool.submit(request, path) for _ in range(4) for path in paths]
for future in as_completed(futures):
path, elapsed, status = future.result()
samples[path].append(elapsed)
if status != 200:
errors.append({"path": path, "status": status, "elapsed_ms": elapsed})
result = {
path: {
"count": len(values),
"p50_ms": percentile(values, 0.5),
"p95_ms": percentile(values, 0.95),
"max_ms": max(values, default=0),
}
for path, values in samples.items()
}
all_values = [value for values in samples.values() for value in values]
result["summary"] = {
"requests": len(all_values),
"concurrency": 4,
"errors": errors,
"p50_ms": percentile(all_values, 0.5),
"p95_ms": percentile(all_values, 0.95),
"max_ms": max(all_values, default=0),
}
if errors or result["summary"]["p95_ms"] > 5000 or result["summary"]["max_ms"] > 10000:
raise AssertionError(f"live performance sample exceeded its guardrail: {result['summary']}")
return result
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)
args = parser.parse_args()
state = json.loads((args.run_dir / "state.json").read_text())
run_id = state["run_id"]
source_id = state["source_id"]
base_url = args.base_url.rstrip("/")
headers = {}
if args.token_file is not None:
headers["Authorization"] = f"Bearer {args.token_file.read_text().strip()}"
client = httpx.Client(
base_url=base_url,
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
status = call(client, "GET", "/api/v1/status").json()
if status.get("version") != "0.5.45" or status.get("access_mode") != "direct":
raise AssertionError(f"unexpected live status: {status}")
unauthenticated = httpx.get(base_url + "/api/v1/profile", timeout=10)
invalid_token = httpx.get(
base_url + "/api/v1/profile",
headers={"Authorization": "Bearer invalid-final-acceptance-token"},
timeout=10,
)
if (unauthenticated.status_code, invalid_token.status_code) != (401, 401):
raise AssertionError("profile endpoint did not reject missing and invalid credentials")
page_one = call(client, "GET", "/api/v1/jobs?page=1&page_size=10").json()
page_two = call(client, "GET", "/api/v1/jobs?page=2&page_size=10").json()
first_ids = {item["id"] for item in page_one["items"]}
second_ids = {item["id"] for item in page_two["items"]}
if first_ids & second_ids or page_one["page"] != 1 or page_two["page"] != 2:
raise AssertionError("background job pagination returned overlapping or incorrect pages")
lane_totals = {
lane: call(client, "GET", f"/api/v1/jobs?page=1&page_size=10&lane={lane}").json()["total"]
for lane in ("ai", "transfer", "download", "scan")
}
if sum(lane_totals.values()) != page_one["total"]:
raise AssertionError(f"job lane totals do not match all jobs: {lane_totals} != {page_one['total']}")
cancel_payload = b"cancelled upload acceptance payload"
cancel_name = f"{run_id}-cancel-before-complete.mp4"
cancel_upload = call(
client,
"POST",
"/api/v1/uploads",
json={
"source_id": source_id,
"relative_path": "ingest",
"filename": cancel_name,
"size_bytes": len(cancel_payload),
"sha256": hashlib.sha256(cancel_payload).hexdigest(),
"conflict": "skip",
},
).json()
cancelled = client.delete(f"/api/v1/uploads/{cancel_upload['id']}")
if cancelled.status_code != 204 or find_video(client, source_id, cancel_name):
raise AssertionError("cancelled receiving upload leaked into the video catalog")
invalid_payload = b"ImageFind final acceptance: intentionally invalid mp4\n"
invalid_name = f"{run_id}-manual-retry-invalid.mp4"
before_job_ids = {item["id"] for item in call(client, "GET", "/api/v1/jobs?limit=500").json()}
invalid_upload = call(
client,
"POST",
"/api/v1/uploads",
json={
"source_id": source_id,
"relative_path": "ingest",
"filename": invalid_name,
"title": "后台任务手动重试验收",
"size_bytes": len(invalid_payload),
"sha256": hashlib.sha256(invalid_payload).hexdigest(),
"conflict": "skip",
},
).json()
call(
client,
"PUT",
f"/api/v1/uploads/{invalid_upload['id']}/chunks/0",
content=invalid_payload,
headers={"X-Chunk-SHA256": hashlib.sha256(invalid_payload).hexdigest()},
)
call(client, "POST", f"/api/v1/uploads/{invalid_upload['id']}/complete")
upload_result = wait_upload(client, invalid_upload["id"])
if upload_result["status"] != "completed":
raise AssertionError(f"invalid fixture did not reach catalog processing: {upload_result}")
deadline = time.monotonic() + 180
failed_job: dict[str, Any] | None = None
invalid_video: dict[str, Any] | None = None
while time.monotonic() < deadline:
invalid_video = find_video(client, source_id, invalid_name)
jobs = call(client, "GET", "/api/v1/jobs?limit=500").json()
failed_job = next(
(
job
for job in jobs
if job["id"] not in before_job_ids
and job["kind"] == "index_video"
and job["status"] == "failed"
),
None,
)
if invalid_video and failed_job:
break
time.sleep(1)
if not invalid_video or not failed_job:
raise AssertionError("invalid video did not produce a run-owned failed index task")
retry = call(client, "POST", f"/api/v1/jobs/{failed_job['id']}/retry", json={}).json()
if retry.get("retried_from") != failed_job["id"] or retry.get("job_id") == failed_job["id"]:
raise AssertionError(f"manual job retry response mismatch: {retry}")
retried_job = wait_job(client, retry["job_id"])
if retried_job["status"] != "failed" or int(retried_job.get("attempts") or 0) < 1:
raise AssertionError(f"retried invalid index job did not execute: {retried_job}")
deleted = call(client, "DELETE", f"/api/v1/videos/{invalid_video['id']}?delete_source=true").json()
if deleted.get("trash_id"):
call(client, "DELETE", f"/api/v1/trash/{deleted['trash_id']}")
resources = call(client, "GET", "/api/v1/system/resources").json()
if resources.get("database", {}).get("engine") != "postgresql":
raise AssertionError(f"unexpected database resource status: {resources.get('database')}")
models = call(client, "GET", "/api/v1/models").json()
if not all(models.get("operational_components", {}).get(name) for name in ("visual", "ocr", "faces", "audio")):
raise AssertionError(f"one or more AI components are not operational: {models.get('operational_components')}")
downloads = call(client, "GET", "/api/v1/downloads/runtime").json()
performance = performance_sample(base_url, client)
report = {
"status": status,
"authentication": {"missing": 401, "invalid": 401},
"jobs": {
"total": page_one["total"],
"pages": page_one["pages"],
"page_size": page_one["page_size"],
"pages_disjoint": True,
"lane_totals": lane_totals,
"failed_job": failed_job,
"retry": retry,
"retried_job": retried_job,
},
"cancelled_upload": {"id": cancel_upload["id"], "hidden_from_videos": True},
"database": resources.get("database"),
"resource_lanes": resources.get("lanes"),
"aria2": downloads,
"models": models.get("operational_components"),
"performance": performance,
}
(args.run_dir / "final-gaps-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()