125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from common import parse_datetime, runtime_instance_id, update_jsonl, utc_now
|
|
|
|
|
|
DEFAULT_CLAIMS_PATH = Path(".modelhub_state/submission_claims.jsonl")
|
|
|
|
|
|
def candidate_key(candidate: dict[str, Any]) -> str:
|
|
model_id = candidate.get("repoId") or candidate.get("modelAddress") or "unknown"
|
|
target_gpu = candidate.get("targetGpu") or "unknown"
|
|
return f"{model_id}|{target_gpu}"
|
|
|
|
|
|
def diversify_candidates(
|
|
candidates: list[dict[str, Any]],
|
|
*,
|
|
instance_id: str,
|
|
) -> list[dict[str, Any]]:
|
|
"""Give concurrent agents different deterministic candidate orders."""
|
|
|
|
def sort_key(candidate: dict[str, Any]) -> str:
|
|
value = f"{instance_id}|{candidate_key(candidate)}"
|
|
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
|
|
|
return sorted(candidates, key=sort_key)
|
|
|
|
|
|
class SubmissionClaimStore:
|
|
"""Small filesystem-backed lease store for local multi-process deduplication."""
|
|
|
|
def __init__(
|
|
self,
|
|
path: Path | str = DEFAULT_CLAIMS_PATH,
|
|
*,
|
|
owner_id: str | None = None,
|
|
claim_ttl_seconds: int = 600,
|
|
submitted_ttl_seconds: int = 24 * 60 * 60,
|
|
) -> None:
|
|
self.path = Path(path)
|
|
self.owner_id = owner_id or runtime_instance_id()
|
|
self.claim_ttl_seconds = max(30, int(claim_ttl_seconds))
|
|
self.submitted_ttl_seconds = max(self.claim_ttl_seconds, int(submitted_ttl_seconds))
|
|
|
|
def claim(
|
|
self,
|
|
candidates: list[dict[str, Any]],
|
|
*,
|
|
limit: int,
|
|
) -> list[dict[str, Any]]:
|
|
if limit <= 0 or not candidates:
|
|
return []
|
|
|
|
selected: list[dict[str, Any]] = []
|
|
now = utc_now()
|
|
expires_at = now + timedelta(seconds=self.claim_ttl_seconds)
|
|
|
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
active = [record for record in records if not self._is_expired(record, now)]
|
|
claimed_keys = {str(record.get("key")) for record in active if record.get("key")}
|
|
for candidate in candidates:
|
|
key = candidate_key(candidate)
|
|
if key in claimed_keys:
|
|
continue
|
|
active.append(
|
|
{
|
|
"key": key,
|
|
"ownerId": self.owner_id,
|
|
"state": "claimed",
|
|
"claimedAt": now.isoformat(),
|
|
"expiresAt": expires_at.isoformat(),
|
|
}
|
|
)
|
|
claimed_keys.add(key)
|
|
selected.append(candidate)
|
|
if len(selected) >= limit:
|
|
break
|
|
return active
|
|
|
|
update_jsonl(self.path, update)
|
|
return selected
|
|
|
|
def mark_submitted(self, candidates: list[dict[str, Any]]) -> None:
|
|
keys = {candidate_key(candidate) for candidate in candidates}
|
|
if not keys:
|
|
return
|
|
now = utc_now()
|
|
expires_at = now + timedelta(seconds=self.submitted_ttl_seconds)
|
|
|
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
active = [record for record in records if not self._is_expired(record, now)]
|
|
for record in active:
|
|
if record.get("ownerId") == self.owner_id and record.get("key") in keys:
|
|
record["state"] = "submitted"
|
|
record["expiresAt"] = expires_at.isoformat()
|
|
return active
|
|
|
|
update_jsonl(self.path, update)
|
|
|
|
def release(self, candidates: list[dict[str, Any]]) -> None:
|
|
keys = {candidate_key(candidate) for candidate in candidates}
|
|
if not keys:
|
|
return
|
|
now = utc_now()
|
|
|
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
return [
|
|
record
|
|
for record in records
|
|
if not self._is_expired(record, now)
|
|
and not (record.get("ownerId") == self.owner_id and record.get("key") in keys)
|
|
]
|
|
|
|
update_jsonl(self.path, update)
|
|
|
|
@staticmethod
|
|
def _is_expired(record: dict[str, Any], now) -> bool: # noqa: ANN001
|
|
expires_at = parse_datetime(record.get("expiresAt"))
|
|
return expires_at is None or expires_at <= now
|