253 lines
15 KiB
Python
253 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""Coordinate four advisory built-in imagegen submissions across HD workers.
|
|
|
|
No API call is started or terminated by this tool. Unstarted reservations last
|
|
thirty minutes. Started and legacy-unknown requests require an explicit finish.
|
|
Workers without a lease should prepare/review other jobs, then try again.
|
|
"""
|
|
import argparse
|
|
import collections
|
|
import datetime
|
|
import fcntl
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
CAPACITY = 4
|
|
LEASE_SECONDS = 30 * 60
|
|
EXPIRY_GRACE_SECONDS = 2 * 60
|
|
RETRY_WINDOW_SECONDS = 30 * 60
|
|
MAX_JOB_CONNECTION_FAILURES = 3
|
|
|
|
|
|
def utc(seconds):
|
|
return datetime.datetime.fromtimestamp(seconds, datetime.timezone.utc).isoformat(timespec="seconds")
|
|
|
|
|
|
def sanitized_error(message):
|
|
"""Preserve the service text apart from credentials or embedded image data."""
|
|
message = re.sub(r"data:image/[^;,\s]+;base64,[A-Za-z0-9+/=]+", "<image-data-redacted>", message)
|
|
message = re.sub(r"\bsk-[A-Za-z0-9_-]{16,}\b", "<api-key-redacted>", message)
|
|
message = re.sub(r"(?i)(\bBearer\s+)[A-Za-z0-9._~-]+", r"\1<credential-redacted>", message)
|
|
message = re.sub(r"(?i)([?&](?:api_key|apikey|access_token|token|signature)=)[^&\s]+", r"\1<credential-redacted>", message)
|
|
return message
|
|
|
|
|
|
def write_atomic(path, value):
|
|
temporary = path.with_name(path.name + f".{os.getpid()}.tmp")
|
|
temporary.write_text(json.dumps(value, indent=2) + "\n")
|
|
temporary.replace(path)
|
|
|
|
|
|
def status(directory, now=None):
|
|
"""Read lease ages and recent outcomes without expiring or changing leases."""
|
|
now = time.time() if now is None else now
|
|
path = directory / "state.json"
|
|
if not path.is_file():
|
|
return {"action": "status", "capacity": CAPACITY, "active_reservations": 0,
|
|
"leases": [], "recent_outcomes": {}}
|
|
with (directory / "gate.lock").open("a") as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_SH)
|
|
state = json.loads(path.read_text())
|
|
if state.get("format_version") != 1:
|
|
raise ValueError("unsupported gate state; do not erase active reservations")
|
|
leases = [{"lease_id": identity, "owner": lease["owner"],
|
|
"job_id": lease["job_id"],
|
|
"phase": lease.get("phase", "legacy_unknown"),
|
|
"age_seconds": max(0, int(now - lease["created_epoch"])),
|
|
"expires_at": utc(lease["expires_epoch"]),
|
|
"renewal_needed": now >= lease["expires_epoch"],
|
|
"stale_reservation": lease.get("phase") == "reserved" and now > lease["expires_epoch"] + EXPIRY_GRACE_SECONDS,
|
|
"owner_status_needed": lease.get("phase", "legacy_unknown") == "legacy_unknown"
|
|
or (lease.get("phase") == "started" and now - lease["created_epoch"] >= LEASE_SECONDS)}
|
|
for identity, lease in state["leases"].items()]
|
|
recent = []
|
|
log = directory / "events.jsonl"
|
|
if log.is_file():
|
|
for line in log.read_text().splitlines():
|
|
event = json.loads(line)
|
|
stamp = datetime.datetime.fromisoformat(event["timestamp_utc"]).timestamp()
|
|
if stamp >= now - RETRY_WINDOW_SECONDS:
|
|
recent.append(event)
|
|
outcomes = collections.Counter(event["outcome"] for event in recent
|
|
if event["event"] == "report")
|
|
acquired = collections.Counter(event["owner"] for event in recent
|
|
if event["event"] == "acquired")
|
|
return {"action": "status", "timestamp_utc": utc(now), "capacity": CAPACITY,
|
|
"admissions_paused": state.get("admissions_paused", False),
|
|
"pause_reason": state.get("pause_reason"),
|
|
"active_reservations": len(leases), "leases": leases,
|
|
"cooldowns": [{"owner": owner, "remaining_seconds": int(row["cooldown_until"] - now)}
|
|
for owner, row in state["owners"].items()
|
|
if row.get("cooldown_until", 0) > now],
|
|
"recent_outcomes": dict(outcomes), "recent_acquisitions_by_owner": dict(acquired),
|
|
"last_event": recent[-1] if recent else None,
|
|
"fairness": "One lease per owner; eligible callers acquire the first available slot. No FIFO wait queue.",
|
|
"limitation": "A reservation age does not prove the imagegen tool is running; ask its owner before changing an active request."}
|
|
|
|
|
|
def operate(directory, action, owner, job=None, lease_id=None, outcome=None, error=None, now=None):
|
|
now = time.time() if now is None else now
|
|
if not owner or len(owner) > 128 or any(char in owner for char in "\r\n\0"):
|
|
raise ValueError("invalid owner identity")
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
with (directory / "gate.lock").open("a") as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX)
|
|
path = directory / "state.json"
|
|
state = json.loads(path.read_text()) if path.is_file() else {"format_version": 1, "leases": {}, "owners": {}, "job_failures": {}}
|
|
if state.get("format_version") != 1:
|
|
raise ValueError("unsupported gate state; do not erase active reservations")
|
|
events = []
|
|
for identity, lease in list(state["leases"].items()):
|
|
if (lease.get("phase") == "reserved" and identity != lease_id
|
|
and now > lease["expires_epoch"] + EXPIRY_GRACE_SECONDS):
|
|
del state["leases"][identity]
|
|
events.append({"event": "expired_stale_reservation", "lease_id": identity, "owner": lease["owner"], "job_id": lease["job_id"], "previous_expiry": utc(lease["expires_epoch"]), "note": "Reservation expiry only; no running imagegen call was interrupted."})
|
|
response = {"action": action, "owner": owner, "capacity": CAPACITY}
|
|
if action in ["pause", "resume"]:
|
|
state["admissions_paused"] = action == "pause"
|
|
state["pause_reason"] = job if action == "pause" else None
|
|
response.update(changed=True, admissions_paused=state["admissions_paused"])
|
|
events.append({"event": action, "owner": owner, "reason": state["pause_reason"],
|
|
"note": "Admission policy only; existing tool requests are not interrupted."})
|
|
elif action == "acquire":
|
|
if not job or not re.fullmatch(r"[A-Za-z0-9_.:-]{1,128}", job):
|
|
raise ValueError("valid job ID required")
|
|
if state.get("admissions_paused"):
|
|
response.update(acquired=False, reason="admissions_paused", pause_reason=state.get("pause_reason"))
|
|
elif any(lease["owner"] == owner for lease in state["leases"].values()):
|
|
response.update(acquired=False, reason="owner_already_has_inflight_reservation")
|
|
elif any(lease["job_id"] == job for lease in state["leases"].values()):
|
|
response.update(acquired=False, reason="job_already_reserved")
|
|
else:
|
|
owner_state = state["owners"].get(owner, {})
|
|
failure_key = owner + "|" + job
|
|
failures = [stamp for stamp in state["job_failures"].get(failure_key, []) if now - stamp < RETRY_WINDOW_SECONDS]
|
|
state["job_failures"][failure_key] = failures
|
|
if len(failures) >= MAX_JOB_CONNECTION_FAILURES:
|
|
response.update(acquired=False, reason="job_connection_retry_batch_exhausted", retry_after_seconds=max(1, int(failures[0] + RETRY_WINDOW_SECONDS - now)))
|
|
elif owner_state.get("cooldown_until", 0) > now:
|
|
response.update(acquired=False, reason="connection_failure_cooldown", retry_after_seconds=max(1, int(owner_state["cooldown_until"] - now)))
|
|
elif len(state["leases"]) >= CAPACITY:
|
|
response.update(acquired=False, reason="all_four_submission_slots_reserved")
|
|
else:
|
|
identity = str(uuid.uuid4())
|
|
lease = {"owner": owner, "job_id": job, "created_epoch": now, "expires_epoch": now + LEASE_SECONDS, "phase": "reserved"}
|
|
state["leases"][identity] = lease
|
|
response.update(acquired=True, lease_id=identity, expires_at=utc(lease["expires_epoch"]), renewal_command=f"renew --owner {owner} --lease {identity}")
|
|
events.append({"event": "acquired", "lease_id": identity, "owner": owner, "job_id": job, "expires_at": utc(lease["expires_epoch"])})
|
|
elif action == "restore":
|
|
if lease_id in state["leases"]:
|
|
if state["leases"][lease_id]["owner"] != owner:
|
|
raise ValueError("reservation belongs to another worker")
|
|
response.update(changed=False, reason="reservation_already_present")
|
|
else:
|
|
history = directory / "events.jsonl"
|
|
source = None
|
|
for line in history.read_text().splitlines() if history.is_file() else []:
|
|
event = json.loads(line)
|
|
if event.get("lease_id") == lease_id:
|
|
if event.get("event") == "acquired":
|
|
source = event
|
|
elif event.get("event") in ["report", "release"]:
|
|
source = None
|
|
if not source or source["owner"] != owner:
|
|
raise ValueError("restore requires an unfinished recorded acquisition belonging to this owner")
|
|
created = datetime.datetime.fromisoformat(source["timestamp_utc"]).timestamp()
|
|
state["leases"][lease_id] = {"owner": owner, "job_id": source["job_id"],
|
|
"created_epoch": created, "expires_epoch": now + LEASE_SECONDS,
|
|
"phase": "legacy_unknown"}
|
|
response.update(changed=True, reason="unverified_expired_request_restored")
|
|
events.append({"event": "restored_unknown_request", "lease_id": lease_id,
|
|
"owner": owner, "job_id": source["job_id"],
|
|
"note": "Retain active-call accounting until the owner confirms completion. No tool was started or interrupted."})
|
|
else:
|
|
lease = state["leases"].get(lease_id)
|
|
if lease is None:
|
|
response.update(changed=False, reason="reservation_missing_or_expired")
|
|
elif lease["owner"] != owner:
|
|
raise ValueError("reservation belongs to another worker")
|
|
elif action == "renew":
|
|
lease["expires_epoch"] = now + LEASE_SECONDS
|
|
response.update(changed=True, expires_at=utc(lease["expires_epoch"]))
|
|
events.append({"event": "renewed", "lease_id": lease_id, "owner": owner, "job_id": lease["job_id"], "expires_at": utc(lease["expires_epoch"])})
|
|
elif action == "start":
|
|
if lease.get("phase") == "started":
|
|
response.update(changed=False, reason="request_already_started")
|
|
elif lease.get("phase") != "reserved":
|
|
response.update(changed=False, reason="legacy_request_requires_recovery_or_finish")
|
|
elif len(state["leases"]) > CAPACITY:
|
|
response.update(changed=False, reason="active_request_accounting_over_capacity")
|
|
elif state.get("admissions_paused"):
|
|
response.update(changed=False, reason="admissions_paused")
|
|
else:
|
|
lease.update(phase="started", started_epoch=now, expires_epoch=now + LEASE_SECONDS)
|
|
response.update(changed=True, started_at=utc(now))
|
|
events.append({"event": "started", "lease_id": lease_id, "owner": owner,
|
|
"job_id": lease["job_id"],
|
|
"note": "Owner is submitting the built-in imagegen request; this slot requires an explicit report or release."})
|
|
else:
|
|
del state["leases"][lease_id]
|
|
event = {"event": action, "lease_id": lease_id, "owner": owner, "job_id": lease["job_id"]}
|
|
response.update(changed=True)
|
|
if action == "report":
|
|
event["outcome"] = outcome
|
|
if error is not None:
|
|
event["sanitized_error"] = sanitized_error(error)
|
|
owner_state = state["owners"].setdefault(owner, {})
|
|
failure_key = owner + "|" + lease["job_id"]
|
|
if outcome == "connection_failure":
|
|
failures = [stamp for stamp in state["job_failures"].get(failure_key, []) if now - stamp < RETRY_WINDOW_SECONDS]
|
|
failures.append(now)
|
|
state["job_failures"][failure_key] = failures
|
|
streak = min(owner_state.get("consecutive_connection_failures", 0) + 1, 10)
|
|
cooldown = min(60 * 2 ** (streak - 1), 600)
|
|
owner_state.update(consecutive_connection_failures=streak, cooldown_until=now + cooldown)
|
|
response.update(retry_after_seconds=cooldown, job_connection_failures_in_batch=len(failures), retry_batch_exhausted=len(failures) >= MAX_JOB_CONNECTION_FAILURES)
|
|
elif outcome == "success":
|
|
state["owners"][owner] = {"consecutive_connection_failures": 0, "cooldown_until": 0}
|
|
state["job_failures"].pop(failure_key, None)
|
|
events.append(event)
|
|
response["active_reservations"] = len(state["leases"])
|
|
response["prepare_other_jobs_when_waiting"] = True
|
|
write_atomic(path, state)
|
|
with (directory / "events.jsonl").open("a") as log:
|
|
for event in events:
|
|
log.write(json.dumps({"timestamp_utc": utc(now), **event}, ensure_ascii=False) + "\n")
|
|
return response
|
|
|
|
|
|
def main():
|
|
project = Path(__file__).resolve().parents[1]
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("action", choices=["acquire", "start", "renew", "release", "report", "restore", "pause", "resume", "status"])
|
|
parser.add_argument("--owner")
|
|
parser.add_argument("--job")
|
|
parser.add_argument("--lease")
|
|
parser.add_argument("--outcome", choices=["success", "connection_failure", "failure"])
|
|
parser.add_argument("--error-file", type=Path, help="UTF-8 service error; credential/image data is redacted, other text preserved")
|
|
parser.add_argument("--directory", type=Path, default=project.parent / "data/hd/jobs/.imagegen-gate")
|
|
args = parser.parse_args()
|
|
if args.action == "status":
|
|
print(json.dumps(status(args.directory), indent=2))
|
|
return
|
|
if not args.owner:
|
|
parser.error("owner required for lease changes")
|
|
if args.action not in ["acquire", "pause", "resume"] and not args.lease:
|
|
parser.error("lease ID required")
|
|
if args.action == "report" and not args.outcome:
|
|
parser.error("report needs an actual outcome")
|
|
error = None
|
|
if args.error_file:
|
|
if args.error_file.stat().st_size > 64 * 1024:
|
|
parser.error("error file exceeds64KiB; remove embedded payloads before logging")
|
|
error = args.error_file.read_text()
|
|
print(json.dumps(operate(args.directory, args.action, args.owner, args.job, args.lease, args.outcome, error), indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|