68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from common import read_jsonl, update_jsonl, utc_now
|
|
|
|
|
|
DEFAULT_SUBMISSION_EXCLUSIONS_PATH = Path(".modelhub_state/submission_exclusions.jsonl")
|
|
|
|
|
|
def exclusion_key(model_id: str, target_gpu: str) -> str:
|
|
return f"{model_id}|{target_gpu}"
|
|
|
|
|
|
class SubmissionExclusionStore:
|
|
"""Persistent model/GPU exclusions for non-retryable platform rejections."""
|
|
|
|
def __init__(self, path: Path | str = DEFAULT_SUBMISSION_EXCLUSIONS_PATH) -> None:
|
|
self.path = Path(path)
|
|
self._lock = threading.Lock()
|
|
self._blocked = {
|
|
exclusion_key(str(record.get("modelId")), str(record.get("targetGpu")))
|
|
for record in read_jsonl(self.path)
|
|
if record.get("modelId") and record.get("targetGpu")
|
|
}
|
|
|
|
def is_blocked(self, model_id: str, target_gpu: str) -> bool:
|
|
with self._lock:
|
|
return exclusion_key(model_id, target_gpu) in self._blocked
|
|
|
|
def block(self, model_id: str, target_gpu: str, *, reason: str) -> None:
|
|
if not model_id or not target_gpu:
|
|
return
|
|
key = exclusion_key(model_id, target_gpu)
|
|
now = utc_now().isoformat()
|
|
|
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
existing = next(
|
|
(
|
|
record
|
|
for record in records
|
|
if exclusion_key(str(record.get("modelId")), str(record.get("targetGpu"))) == key
|
|
),
|
|
None,
|
|
)
|
|
if existing is None:
|
|
records.append(
|
|
{
|
|
"modelId": model_id,
|
|
"targetGpu": target_gpu,
|
|
"reason": reason,
|
|
"firstSeenAt": now,
|
|
"lastSeenAt": now,
|
|
"occurrences": 1,
|
|
}
|
|
)
|
|
else:
|
|
existing["reason"] = reason
|
|
existing["lastSeenAt"] = now
|
|
existing["occurrences"] = max(1, int(existing.get("occurrences") or 1)) + 1
|
|
return records
|
|
|
|
update_jsonl(self.path, update)
|
|
with self._lock:
|
|
self._blocked.add(key)
|