51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from check_inference_endpoint import check_endpoint
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_REPORT = ROOT / "reports" / "inference-endpoint-wait.json"
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Wait for an OpenAI-compatible inference endpoint.")
|
|
parser.add_argument("--base-url", required=True)
|
|
parser.add_argument("--expected-model")
|
|
parser.add_argument("--timeout", type=int, default=900)
|
|
parser.add_argument("--interval", type=int, default=15)
|
|
parser.add_argument("--request-timeout", type=int, default=10)
|
|
parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
|
|
parser.add_argument("--print", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
deadline = time.monotonic() + args.timeout
|
|
attempts = []
|
|
while True:
|
|
result = check_endpoint(args.base_url, args.expected_model, args.request_timeout)
|
|
result["attempt"] = len(attempts) + 1
|
|
attempts.append(result)
|
|
if result["status"] == "ok":
|
|
report = {"status": "ok", "attempts": attempts}
|
|
break
|
|
if time.monotonic() >= deadline:
|
|
report = {"status": "failed", "attempts": attempts}
|
|
break
|
|
time.sleep(args.interval)
|
|
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
if args.print:
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
else:
|
|
print(f"Endpoint wait status: {report['status']} ({args.base_url})")
|
|
return 0 if report["status"] == "ok" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|