739 lines
31 KiB
Python
739 lines
31 KiB
Python
from __future__ import annotations
|
|
|
|
import fcntl
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import threading
|
|
import uuid
|
|
import re
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
from dulwich import porcelain
|
|
from dulwich.repo import Repo
|
|
|
|
from common import read_json, read_jsonl, write_json, write_jsonl
|
|
from runner_common import load_key_files
|
|
from version import AGENT_VERSION
|
|
|
|
|
|
STATE_SCHEMA_VERSION = 1
|
|
DEFAULT_REMOTE = "https://dev.modelhub.org.cn/CoolBoy/submmit.git"
|
|
DEFAULT_BRANCH = "agent-state"
|
|
DEFAULT_BATCH_SIZE = 100
|
|
DEFAULT_HISTORY_DEPTH = 200
|
|
DEFAULT_RECENT_TERMINAL_INTENTS = 300
|
|
DEFAULT_LEDGER_RECORDS = 500
|
|
DEFAULT_CRASH_RECORDS = 50
|
|
|
|
STATE_OUTCOME_FIELDS = {
|
|
"taskId", "modelId", "targetGpu", "framework", "taskType", "submitTime",
|
|
"lastSyncTime", "status", "verifyResult", "outcome", "failReason",
|
|
"modelProfile", "failureCode", "failureCategory", "failureScope",
|
|
"failureAction", "failureDeterministic", "failureNeedsLlm",
|
|
"failureClassificationReason", "failureObservedGpuMemoryGiB",
|
|
"failureUnsupportedArchitectures", "failureUnsupportedModelTypes",
|
|
"failureDetectedFramework", "failureEnrichmentAttempts",
|
|
"failureEnrichmentError", "platformFailure", "policyCancelled",
|
|
"policyCancellationReasons", "policyCancelledAt",
|
|
"policyCancellationResolvedAsSuccess",
|
|
}
|
|
|
|
# Only these runtime files may cross the trust boundary into the state branch.
|
|
# Credentials, raw stdout, downloaded archives and run directories are excluded.
|
|
STATE_ALLOWLIST = (
|
|
".modelhub_state/account_capacity.json",
|
|
".modelhub_state/architecture_compatibility_blacklist.json",
|
|
".modelhub_state/gpu_strategy.json",
|
|
".modelhub_state/market_intelligence.json",
|
|
".modelhub_state/official_capabilities.json",
|
|
".modelhub_state/outcome_checkpoint.json",
|
|
".modelhub_state/queue_cleanup_latest.json",
|
|
".modelhub_state/recent_outcomes.jsonl",
|
|
".modelhub_state/recovery_active_tasks.jsonl",
|
|
".modelhub_state/recovery_intents.jsonl",
|
|
".modelhub_state/routing_intelligence.json",
|
|
".modelhub_state/submission_exclusions.jsonl",
|
|
".modelhub_state/worker_crashes.jsonl",
|
|
"ledger/submissions.jsonl",
|
|
"outcomes/submissions.jsonl",
|
|
)
|
|
|
|
|
|
class StateSyncError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _utc_now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _safe_text(value: Any, limit: int = 2000) -> str:
|
|
text = str(value or "")
|
|
for marker in ("xc-token", "authorization", "password", "token="):
|
|
if marker in text.lower():
|
|
return "<redacted>"
|
|
return text[:limit]
|
|
|
|
|
|
def _safe_config_vector(config: str) -> dict[str, Any]:
|
|
"""Extract only bounded tuning values; never persist the original config."""
|
|
vector: dict[str, Any] = {}
|
|
patterns = {
|
|
"gpuNum": r"\bgpu_num\s*:\s*['\"]?(\d+)",
|
|
"tensorParallel": r"(?:--tensor-parallel-size|-tp)\s*[, ]?\s*['\"]?(\d+)",
|
|
"maxModelLen": r"(?:--max-model-len|max_model_len|max_seq_len)\s*[: ,]+\s*['\"]?(\d+)",
|
|
"gpuMemoryUtilization": r"(?:--gpu-memory-utilization|gpu_memory_utilization)\s*[: ,]+\s*['\"]?([0-9.]+)",
|
|
}
|
|
for key, pattern in patterns.items():
|
|
values = re.findall(pattern, config, flags=re.IGNORECASE)
|
|
if not values:
|
|
continue
|
|
try:
|
|
parsed = float(values[-1]) if "." in values[-1] else int(values[-1])
|
|
except ValueError:
|
|
continue
|
|
vector[key] = parsed
|
|
for key, pattern in {
|
|
"dtype": r"(?:--dtype|dtype)\s*[: ,]+\s*['\"]?([A-Za-z0-9_-]+)",
|
|
"quantization": r"(?:--quantization|quantization)\s*[: ,]+\s*['\"]?([A-Za-z0-9_-]+)",
|
|
"loadFormat": r"(?:--load-format|load_format)\s*[: ,]+\s*['\"]?([A-Za-z0-9_-]+)",
|
|
}.items():
|
|
values = re.findall(pattern, config, flags=re.IGNORECASE)
|
|
if values:
|
|
vector[key] = values[-1][:64]
|
|
return vector
|
|
|
|
|
|
def load_state_git_credentials(*dotenv_paths: Path) -> dict[str, str]:
|
|
paths = dotenv_paths or (Path(".env"), Path(__file__).resolve().parent.parent / ".env")
|
|
values = load_key_files(*paths)
|
|
|
|
def choose(env_name: str, dotenv_name: str) -> str:
|
|
raw = os.getenv(env_name) or values.get(dotenv_name) or ""
|
|
return str(raw).strip().strip('"').strip("'")
|
|
|
|
return {
|
|
"username": choose("MODELHUB_GIT_USERNAME", "modelhub_user_name"),
|
|
"email": choose("MODELHUB_GIT_EMAIL", "modelhub_user_email"),
|
|
"password": choose("MODELHUB_GIT_PASSWORD", "modelhub_user_password"),
|
|
}
|
|
|
|
|
|
def state_git_credentials_present(*dotenv_paths: Path) -> bool:
|
|
credentials = load_state_git_credentials(*dotenv_paths)
|
|
return all(credentials.values())
|
|
|
|
|
|
class StateGitSync:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
project_root: Path | str,
|
|
credentials: dict[str, str],
|
|
remote: str = DEFAULT_REMOTE,
|
|
branch: str = DEFAULT_BRANCH,
|
|
batch_size: int = DEFAULT_BATCH_SIZE,
|
|
history_depth: int = DEFAULT_HISTORY_DEPTH,
|
|
log_fn=None,
|
|
) -> None:
|
|
self.project_root = Path(project_root).resolve()
|
|
self.credentials = dict(credentials)
|
|
self.remote = remote
|
|
self.branch = branch
|
|
self.batch_size = max(1, min(100, int(batch_size)))
|
|
self.history_depth = max(2, int(history_depth))
|
|
self.log = log_fn or (lambda message: print(message, flush=True))
|
|
self.writer_id = uuid.uuid4().hex
|
|
self.generation = 0
|
|
self.healthy = False
|
|
self.last_error: str | None = None
|
|
self.last_sync_at: str | None = None
|
|
self._lock_handle = None
|
|
self._mutex = threading.Lock()
|
|
self._workspace: Path | None = None
|
|
self._expected_remote_oid: str | None = None
|
|
self._pending_push_oid: str | None = None
|
|
self._pending_generation: int | None = None
|
|
self._pending_manifest: dict[str, Any] | None = None
|
|
|
|
@property
|
|
def state_dir(self) -> Path:
|
|
return self.project_root / ".modelhub_state"
|
|
|
|
@property
|
|
def intents_path(self) -> Path:
|
|
return self.state_dir / "recovery_intents.jsonl"
|
|
|
|
@property
|
|
def active_tasks_path(self) -> Path:
|
|
return self.state_dir / "recovery_active_tasks.jsonl"
|
|
|
|
@property
|
|
def readiness_path(self) -> Path:
|
|
return self.state_dir / "readiness.json"
|
|
|
|
def acquire_process_lock(self) -> None:
|
|
self.state_dir.mkdir(parents=True, exist_ok=True)
|
|
path = self.state_dir / "state_sync.lock"
|
|
handle = path.open("a+", encoding="utf-8")
|
|
try:
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError as exc:
|
|
handle.close()
|
|
raise StateSyncError("another state-sync writer is already running") from exc
|
|
self._lock_handle = handle
|
|
|
|
def close(self) -> None:
|
|
if self._lock_handle is not None:
|
|
try:
|
|
fcntl.flock(self._lock_handle.fileno(), fcntl.LOCK_UN)
|
|
finally:
|
|
self._lock_handle.close()
|
|
self._lock_handle = None
|
|
if self._workspace is not None:
|
|
shutil.rmtree(self._workspace.parent, ignore_errors=True)
|
|
self._workspace = None
|
|
|
|
def _auth_kwargs(self) -> dict[str, str]:
|
|
if not all(self.credentials.get(key) for key in ("username", "email", "password")):
|
|
raise StateSyncError("missing ModelHub Git username, email or password")
|
|
return {
|
|
"username": self.credentials["username"],
|
|
"password": self.credentials["password"],
|
|
}
|
|
|
|
@property
|
|
def _author(self) -> bytes:
|
|
name = self.credentials.get("username") or "modelhub-agent"
|
|
email = self.credentials.get("email") or "modelhub-agent@localhost"
|
|
return f"{name} <{email}>".encode("utf-8")
|
|
|
|
def _remote_oid(self) -> str | None:
|
|
return self._remote_oid_for(self.branch)
|
|
|
|
def _remote_oid_for(self, branch: str) -> str | None:
|
|
result = porcelain.ls_remote(self.remote, quiet=True, **self._auth_kwargs())
|
|
oid = result.refs.get(f"refs/heads/{branch}".encode("utf-8"))
|
|
return oid.decode("ascii") if oid else None
|
|
|
|
def _create_workspace(self, remote_oid: str | None) -> None:
|
|
parent = Path(tempfile.mkdtemp(prefix="modelhub-agent-state-"))
|
|
workspace = parent / "state"
|
|
if remote_oid:
|
|
porcelain.clone(
|
|
self.remote,
|
|
workspace,
|
|
branch=self.branch,
|
|
depth=self.history_depth,
|
|
checkout=True,
|
|
errstream=io.BytesIO(),
|
|
**self._auth_kwargs(),
|
|
)
|
|
else:
|
|
workspace.mkdir(parents=True)
|
|
repo = porcelain.init(workspace)
|
|
repo.refs.set_symbolic_ref(b"HEAD", f"refs/heads/{self.branch}".encode("utf-8"))
|
|
self._workspace = workspace
|
|
self._expected_remote_oid = remote_oid
|
|
|
|
def _validate_manifest(self, root: Path) -> dict[str, Any]:
|
|
manifest = read_json(root / "manifest.json")
|
|
if not isinstance(manifest, dict) or int(manifest.get("schemaVersion") or 0) != STATE_SCHEMA_VERSION:
|
|
raise StateSyncError("unsupported or missing state manifest")
|
|
checksums = manifest.get("checksums") or {}
|
|
if not isinstance(checksums, dict):
|
|
raise StateSyncError("invalid state manifest checksums")
|
|
for relative, expected in checksums.items():
|
|
if relative not in STATE_ALLOWLIST and not str(relative).startswith("events/"):
|
|
raise StateSyncError(f"state manifest contains a non-allowlisted path: {relative}")
|
|
path = root / str(relative)
|
|
if not path.is_file() or _sha256_file(path) != str(expected):
|
|
raise StateSyncError(f"state checksum mismatch: {relative}")
|
|
return manifest
|
|
|
|
def restore(self) -> bool:
|
|
with self._mutex:
|
|
try:
|
|
remote_oid = self._remote_oid()
|
|
self._create_workspace(remote_oid)
|
|
if remote_oid is not None:
|
|
try:
|
|
manifest = self._validate_manifest(self._workspace)
|
|
except Exception as current_error:
|
|
manifest = None
|
|
repo = Repo(str(self._workspace))
|
|
commits = [entry.commit.id for entry in repo.get_walker(max_entries=self.history_depth)][1:]
|
|
for commit in commits:
|
|
porcelain.reset(repo, "hard", commit)
|
|
try:
|
|
manifest = self._validate_manifest(self._workspace)
|
|
self.log(
|
|
f"[state-recovery] fallback_commit={commit.decode('ascii')[:12]} "
|
|
"reason=checksum_recovery"
|
|
)
|
|
break
|
|
except Exception:
|
|
continue
|
|
if manifest is None:
|
|
raise current_error
|
|
self.generation = max(0, int(manifest.get("generation") or 0))
|
|
for relative in STATE_ALLOWLIST:
|
|
source = self._workspace / relative
|
|
if not source.is_file():
|
|
continue
|
|
destination = self.project_root / relative
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = destination.with_name(f".{destination.name}.restore-{uuid.uuid4().hex}")
|
|
if relative == ".modelhub_state/worker_crashes.jsonl":
|
|
merged: dict[str, dict[str, Any]] = {}
|
|
for row in [*read_jsonl(source), *read_jsonl(destination)]:
|
|
key = json.dumps(row, ensure_ascii=False, sort_keys=True)
|
|
merged[key] = row
|
|
rows = sorted(merged.values(), key=lambda row: str(row.get("at") or ""))[-DEFAULT_CRASH_RECORDS:]
|
|
write_jsonl(temporary, rows)
|
|
else:
|
|
shutil.copy2(source, temporary)
|
|
os.replace(temporary, destination)
|
|
self.log(
|
|
f"[state-recovery] generation={self.generation} source=remote branch={self.branch} status=ok"
|
|
)
|
|
else:
|
|
self.log(f"[state-recovery] source=local_bootstrap branch={self.branch} status=ok")
|
|
self.healthy = True
|
|
self.last_error = None
|
|
return True
|
|
except Exception as exc:
|
|
self.healthy = False
|
|
self.last_error = _safe_text(exc)
|
|
self.log(f"[state-recovery] status=failed reason={self.last_error}")
|
|
return False
|
|
|
|
def retry_restore(self) -> bool:
|
|
if self._workspace is not None:
|
|
shutil.rmtree(self._workspace.parent, ignore_errors=True)
|
|
self._workspace = None
|
|
self._expected_remote_oid = None
|
|
return self.restore()
|
|
|
|
def _compact_intents(self) -> int:
|
|
records = read_jsonl(self.intents_path)
|
|
active_statuses = {"pending", "submitted", "recovered_active"}
|
|
active = [row for row in records if str(row.get("status") or "") in active_statuses]
|
|
terminal = [row for row in records if str(row.get("status") or "") not in active_statuses]
|
|
terminal.sort(
|
|
key=lambda row: str(row.get("completedAt") or row.get("createdAt") or ""),
|
|
reverse=True,
|
|
)
|
|
retained_terminal = terminal[:DEFAULT_RECENT_TERMINAL_INTENTS]
|
|
discarded = terminal[DEFAULT_RECENT_TERMINAL_INTENTS:]
|
|
compact_fields = {
|
|
"intentId", "batchId", "status", "createdAt", "completedAt",
|
|
"repoId", "targetGpu", "taskType", "framework", "configSource",
|
|
"configFingerprint", "safeConfigVector", "taskId",
|
|
}
|
|
compacted = [
|
|
*active,
|
|
*(
|
|
{key: value for key, value in row.items() if key in compact_fields}
|
|
for row in retained_terminal
|
|
),
|
|
]
|
|
if compacted == records:
|
|
return 0
|
|
write_jsonl(self.intents_path, compacted)
|
|
self.log(
|
|
f"[state-compact] mode=decision_state_only intents_discarded={len(discarded)} "
|
|
f"active={len(active)} recent_terminal={len(retained_terminal)}"
|
|
)
|
|
return len(discarded)
|
|
|
|
def _event_files(self) -> list[Path]:
|
|
# The recovery intent WAL already records both transitions. Persisting
|
|
# a second event stream doubled state without improving recovery.
|
|
return []
|
|
|
|
@staticmethod
|
|
def _intent(candidate: dict[str, Any], batch_id: str) -> dict[str, Any]:
|
|
config = str(candidate.get("configParams") or "")
|
|
return {
|
|
"intentId": uuid.uuid4().hex,
|
|
"batchId": batch_id,
|
|
"status": "pending",
|
|
"createdAt": _utc_now().isoformat(),
|
|
"repoId": candidate.get("repoId"),
|
|
"modelAddress": candidate.get("modelAddress"),
|
|
"targetGpu": candidate.get("targetGpu"),
|
|
"taskType": candidate.get("taskType"),
|
|
"framework": candidate.get("framework"),
|
|
"lastModified": candidate.get("lastModified"),
|
|
"configSource": candidate.get("frameworkConfigSource") or "official",
|
|
"configFingerprint": hashlib.sha256(config.encode("utf-8")).hexdigest(),
|
|
"safeConfigVector": _safe_config_vector(config),
|
|
}
|
|
|
|
def begin_batch(self, candidates: list[dict[str, Any]]) -> str | None:
|
|
if not self.healthy:
|
|
return None
|
|
batch_id = uuid.uuid4().hex
|
|
records = read_jsonl(self.intents_path)
|
|
intents = [self._intent(candidate, batch_id) for candidate in candidates]
|
|
records.extend(intents)
|
|
write_jsonl(self.intents_path, records)
|
|
if not self.sync("intent"):
|
|
return None
|
|
return batch_id
|
|
|
|
def finish_batch(self, batch_id: str, results: Iterable[dict[str, Any]]) -> bool:
|
|
records = read_jsonl(self.intents_path)
|
|
by_key = {
|
|
(
|
|
str(item.get("repoId") or ""),
|
|
str(item.get("targetGpu") or ""),
|
|
str(item.get("framework") or ""),
|
|
): item
|
|
for item in records
|
|
if item.get("batchId") == batch_id
|
|
}
|
|
for result in results:
|
|
candidate = result.get("candidate") or {}
|
|
key = (
|
|
str(candidate.get("repoId") or ""),
|
|
str(candidate.get("targetGpu") or ""),
|
|
str(candidate.get("framework") or ""),
|
|
)
|
|
intent = by_key.get(key)
|
|
if intent is None:
|
|
continue
|
|
intent["status"] = str(result.get("outcome") or "unknown")
|
|
intent["completedAt"] = _utc_now().isoformat()
|
|
intent["taskId"] = result.get("taskId")
|
|
intent["reason"] = _safe_text(result.get("reason")) if result.get("reason") else None
|
|
write_jsonl(self.intents_path, records)
|
|
return self.sync("result")
|
|
|
|
def record_active_tasks(self, tasks: list[dict[str, Any]]) -> None:
|
|
sanitized: list[dict[str, Any]] = []
|
|
for task in tasks:
|
|
sanitized.append(
|
|
{
|
|
key: task.get(key)
|
|
for key in (
|
|
"accountKey",
|
|
"taskId",
|
|
"modelId",
|
|
"modelAddress",
|
|
"gpuType",
|
|
"targetGpu",
|
|
"taskType",
|
|
"modelTaskLevel",
|
|
"modelTaskLevelId",
|
|
"framework",
|
|
"configFingerprint",
|
|
"status",
|
|
"createTime",
|
|
"updateTime",
|
|
)
|
|
if task.get(key) is not None
|
|
}
|
|
)
|
|
write_jsonl(self.active_tasks_path, sanitized)
|
|
|
|
def reconcile_active_tasks(self, tasks: list[dict[str, Any]]) -> dict[str, int]:
|
|
intents = read_jsonl(self.intents_path)
|
|
recoverable = [
|
|
item
|
|
for item in intents
|
|
if str(item.get("status") or "") in {"pending", "submitted", "recovered_active"}
|
|
]
|
|
by_route = {
|
|
(
|
|
str(item.get("repoId") or ""),
|
|
str(item.get("targetGpu") or ""),
|
|
): item
|
|
for item in recoverable
|
|
}
|
|
reconciled = 0
|
|
enriched: list[dict[str, Any]] = []
|
|
active_routes: set[tuple[str, str]] = set()
|
|
for task in tasks:
|
|
route = (
|
|
str(task.get("modelId") or task.get("repoId") or ""),
|
|
str(task.get("gpuType") or task.get("targetGpu") or ""),
|
|
)
|
|
active_routes.add(route)
|
|
matched = by_route.get(route)
|
|
copied = dict(task)
|
|
if matched is not None:
|
|
if not copied.get("framework"):
|
|
copied["framework"] = matched.get("framework")
|
|
copied["configFingerprint"] = matched.get("configFingerprint")
|
|
matched["status"] = "recovered_active"
|
|
matched["reconciledAt"] = _utc_now().isoformat()
|
|
if copied.get("taskId") is not None:
|
|
matched["taskId"] = copied.get("taskId")
|
|
reconciled += 1
|
|
enriched.append(copied)
|
|
|
|
now = _utc_now()
|
|
unresolved = 0
|
|
outcome_by_task = {
|
|
str(item.get("taskId")): item
|
|
for item in read_jsonl(self.project_root / "outcomes/submissions.jsonl")
|
|
if item.get("taskId") is not None and item.get("outcome") in {"success", "failed", "policy_cancelled"}
|
|
}
|
|
for intent in recoverable:
|
|
route = (str(intent.get("repoId") or ""), str(intent.get("targetGpu") or ""))
|
|
if route in active_routes:
|
|
continue
|
|
terminal = outcome_by_task.get(str(intent.get("taskId") or ""))
|
|
if terminal is not None:
|
|
intent["status"] = str(terminal.get("outcome") or "terminal")
|
|
intent["completedAt"] = now.isoformat()
|
|
continue
|
|
created = str(intent.get("createdAt") or "")
|
|
try:
|
|
created_at = datetime.fromisoformat(created.replace("Z", "+00:00"))
|
|
except ValueError:
|
|
created_at = now
|
|
if created_at.tzinfo is None:
|
|
created_at = created_at.replace(tzinfo=timezone.utc)
|
|
if intent.get("status") == "pending" and now - created_at >= timedelta(hours=2):
|
|
intent["status"] = "orphan_unconfirmed"
|
|
intent["completedAt"] = now.isoformat()
|
|
elif intent.get("status") == "pending":
|
|
unresolved += 1
|
|
# Terminal intents are compacted by count and reduced to decision
|
|
# fields during snapshot creation. No full historical rows are kept.
|
|
write_jsonl(self.intents_path, intents)
|
|
self.record_active_tasks(enriched)
|
|
return {"active": len(enriched), "reconciled": reconciled, "unresolved": unresolved}
|
|
|
|
def _copy_snapshot(self) -> dict[str, str]:
|
|
assert self._workspace is not None
|
|
self._compact_intents()
|
|
checksums: dict[str, str] = {}
|
|
for relative in STATE_ALLOWLIST:
|
|
source = self.project_root / relative
|
|
destination = self._workspace / relative
|
|
if not source.is_file():
|
|
destination.unlink(missing_ok=True)
|
|
continue
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
if relative == ".modelhub_state/market_intelligence.json":
|
|
payload = read_json(source)
|
|
|
|
def strip_configs(value: Any) -> Any:
|
|
if isinstance(value, dict):
|
|
return {
|
|
key: strip_configs(item)
|
|
for key, item in value.items()
|
|
if key not in {"officialConfig", "configParams"}
|
|
}
|
|
if isinstance(value, list):
|
|
return [strip_configs(item) for item in value]
|
|
return value
|
|
|
|
sanitized_market = strip_configs(payload)
|
|
if isinstance(sanitized_market, dict):
|
|
# A restored snapshot must refresh official configs before routing.
|
|
sanitized_market["frameworkUpdatedAt"] = None
|
|
# Raw community rows are a disposable refresh cache. The
|
|
# aggregated framework/GPU statistics are the durable input.
|
|
sanitized_market["communitySample"] = {}
|
|
write_json(destination, sanitized_market)
|
|
elif relative == ".modelhub_state/official_capabilities.json":
|
|
payload = read_json(source)
|
|
cache = payload.get("modelGpuTaskTypes") if isinstance(payload, dict) else None
|
|
if isinstance(cache, dict) and len(cache) > 750:
|
|
ordered = sorted(
|
|
cache.items(),
|
|
key=lambda pair: str((pair[1] or {}).get("updatedAt") or ""),
|
|
reverse=True,
|
|
)
|
|
payload["modelGpuTaskTypes"] = dict(ordered[:750])
|
|
write_json(destination, payload)
|
|
elif relative in {"outcomes/submissions.jsonl", ".modelhub_state/recent_outcomes.jsonl"}:
|
|
sanitized_outcomes: list[dict[str, Any]] = []
|
|
for row in read_jsonl(source):
|
|
sanitized_outcomes.append(
|
|
{
|
|
key: value
|
|
for key, value in row.items()
|
|
if key in STATE_OUTCOME_FIELDS
|
|
}
|
|
)
|
|
if relative == ".modelhub_state/recent_outcomes.jsonl":
|
|
sanitized_outcomes = sanitized_outcomes[:300]
|
|
write_jsonl(destination, sanitized_outcomes)
|
|
elif relative == "ledger/submissions.jsonl":
|
|
ledger_rows = read_jsonl(source)
|
|
active_ids = {
|
|
str(row.get("taskId"))
|
|
for row in read_jsonl(self.active_tasks_path)
|
|
if row.get("taskId") is not None
|
|
}
|
|
selected_by_task: dict[str, dict[str, Any]] = {}
|
|
anonymous: list[dict[str, Any]] = []
|
|
for row in [
|
|
*(item for item in ledger_rows if str(item.get("taskId") or "") in active_ids),
|
|
*ledger_rows[-DEFAULT_LEDGER_RECORDS:],
|
|
]:
|
|
task_id = str(row.get("taskId") or "")
|
|
if task_id:
|
|
selected_by_task[task_id] = row
|
|
else:
|
|
anonymous.append(row)
|
|
write_jsonl(destination, [*selected_by_task.values(), *anonymous[-20:]])
|
|
elif relative == ".modelhub_state/worker_crashes.jsonl":
|
|
write_jsonl(destination, read_jsonl(source)[-DEFAULT_CRASH_RECORDS:])
|
|
else:
|
|
shutil.copy2(source, destination)
|
|
checksums[relative] = _sha256_file(destination)
|
|
for source in self._event_files():
|
|
relative = f"events/{source.name}"
|
|
destination = self._workspace / relative
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(source, destination)
|
|
checksums[relative] = _sha256_file(destination)
|
|
remote_event_dir = self._workspace / "events"
|
|
if remote_event_dir.exists():
|
|
allowed = {Path(path).name for path in checksums if path.startswith("events/")}
|
|
for path in remote_event_dir.glob("*.jsonl"):
|
|
if path.name not in allowed:
|
|
path.unlink(missing_ok=True)
|
|
return checksums
|
|
|
|
def _push_pending(self, *, phase: str) -> bool:
|
|
if self._pending_push_oid is None:
|
|
return True
|
|
assert self._workspace is not None
|
|
repo = Repo(str(self._workspace))
|
|
remote_oid = self._remote_oid()
|
|
if remote_oid == self._pending_push_oid:
|
|
pushed = True
|
|
elif remote_oid != self._expected_remote_oid:
|
|
raise StateSyncError("state branch changed remotely; refusing to overwrite a newer snapshot")
|
|
else:
|
|
porcelain.push(
|
|
repo,
|
|
self.remote,
|
|
refspecs=f"HEAD:refs/heads/{self.branch}",
|
|
force=True,
|
|
outstream=io.BytesIO(),
|
|
errstream=io.BytesIO(),
|
|
**self._auth_kwargs(),
|
|
)
|
|
pushed = self._remote_oid() == self._pending_push_oid
|
|
if not pushed:
|
|
raise StateSyncError("state branch verification failed after push")
|
|
self._expected_remote_oid = self._pending_push_oid
|
|
self.generation = int(self._pending_generation or self.generation)
|
|
manifest = self._pending_manifest or {}
|
|
self.last_sync_at = manifest.get("updatedAt")
|
|
self._pending_push_oid = None
|
|
self._pending_generation = None
|
|
self._pending_manifest = None
|
|
self.last_error = None
|
|
self.healthy = True
|
|
self.log(f"[state-sync] generation={self.generation} phase={phase} status=ok")
|
|
return True
|
|
|
|
def sync(self, phase: str) -> bool:
|
|
with self._mutex:
|
|
try:
|
|
if self._workspace is None:
|
|
raise StateSyncError("state workspace is not initialized")
|
|
# A failed network push is retried byte-for-byte. Do not create
|
|
# another commit or generation for the same state transition.
|
|
if self._pending_push_oid is not None:
|
|
return self._push_pending(phase=phase)
|
|
previous_manifest: dict[str, Any] = {}
|
|
try:
|
|
previous_manifest = read_json(self._workspace / "manifest.json")
|
|
except (FileNotFoundError, ValueError, TypeError):
|
|
pass
|
|
checksums = self._copy_snapshot()
|
|
if (
|
|
previous_manifest.get("checksums") == checksums
|
|
and previous_manifest.get("agentVersion") == AGENT_VERSION
|
|
):
|
|
self.healthy = True
|
|
self.last_error = None
|
|
return True
|
|
next_generation = self.generation + 1
|
|
manifest = {
|
|
"schemaVersion": STATE_SCHEMA_VERSION,
|
|
"generation": next_generation,
|
|
"updatedAt": _utc_now().isoformat(),
|
|
"agentVersion": AGENT_VERSION,
|
|
"writerId": self.writer_id,
|
|
"phase": phase,
|
|
"checksums": checksums,
|
|
}
|
|
write_json(self._workspace / "manifest.json", manifest)
|
|
repo = Repo(str(self._workspace))
|
|
porcelain.add(repo)
|
|
status = porcelain.status(repo)
|
|
staged = any(status.staged.get(kind) for kind in ("add", "delete", "modify"))
|
|
if not staged:
|
|
self.healthy = True
|
|
return True
|
|
porcelain.commit(
|
|
repo,
|
|
message=f"state: generation {next_generation} ({phase})".encode("utf-8"),
|
|
author=self._author,
|
|
committer=self._author,
|
|
)
|
|
commit_count = sum(1 for _ in repo.get_walker())
|
|
if commit_count > self.history_depth:
|
|
branch_ref = f"refs/heads/{self.branch}".encode("utf-8")
|
|
del repo.refs[branch_ref]
|
|
porcelain.add(repo)
|
|
porcelain.commit(
|
|
repo,
|
|
message=f"state: compacted generation {next_generation}".encode("utf-8"),
|
|
author=self._author,
|
|
committer=self._author,
|
|
)
|
|
local_oid = repo.head().decode("ascii")
|
|
self._pending_push_oid = local_oid
|
|
self._pending_generation = next_generation
|
|
self._pending_manifest = manifest
|
|
return self._push_pending(phase=phase)
|
|
except Exception as exc:
|
|
self.healthy = False
|
|
self.last_error = _safe_text(exc)
|
|
self.log(f"[state-sync] phase={phase} status=failed reason={self.last_error}")
|
|
return False
|
|
|
|
def write_readiness(self, *, ready: bool, reason: str | None = None, extra: dict[str, Any] | None = None) -> None:
|
|
body = {
|
|
"ready": bool(ready),
|
|
"reason": reason,
|
|
"updatedAt": _utc_now().isoformat(),
|
|
"pid": os.getpid(),
|
|
"stateSync": {
|
|
"healthy": self.healthy,
|
|
"generation": self.generation,
|
|
"lastSyncAt": self.last_sync_at,
|
|
"lastError": self.last_error,
|
|
},
|
|
}
|
|
if extra:
|
|
body.update(extra)
|
|
write_json(self.readiness_path, body)
|