Files
submmit/modelhub_submmit_api/state_sync.py
2026-08-21 03:38:06 +08:00

657 lines
27 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 = 20
DEFAULT_RETENTION_DAYS = 30
DEFAULT_HISTORY_DEPTH = 20
# 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/queue_cleanup_latest.json",
".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,
retention_days: int = DEFAULT_RETENTION_DAYS,
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.retention_days = max(1, int(retention_days))
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
@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:
result = porcelain.ls_remote(self.remote, quiet=True, **self._auth_kwargs())
oid = result.refs.get(f"refs/heads/{self.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 ""))[-200:]
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 _event_files(self) -> list[Path]:
event_dir = self.state_dir / "events"
if not event_dir.exists():
return []
cutoff = (_utc_now() - timedelta(days=self.retention_days)).date()
result: list[Path] = []
for path in sorted(event_dir.glob("*.jsonl")):
try:
event_day = datetime.strptime(path.stem, "%Y-%m-%d").date()
except ValueError:
continue
if event_day >= cutoff:
result.append(path)
else:
path.unlink(missing_ok=True)
return result
def _append_event(self, event: dict[str, Any]) -> None:
now = _utc_now()
path = self.state_dir / "events" / f"{now.date().isoformat()}.jsonl"
existing = read_jsonl(path)
sanitized = {key: value for key, value in event.items() if key not in {"configParams", "token", "password"}}
sanitized["at"] = sanitized.get("at") or now.isoformat()
sanitized["eventId"] = sanitized.get("eventId") or uuid.uuid4().hex
if "reason" in sanitized:
sanitized["reason"] = _safe_text(sanitized["reason"])
existing.append(sanitized)
write_jsonl(path, existing)
@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)
for intent in intents:
self._append_event({**intent, "event": "submission_intent"})
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
self._append_event({**intent, "event": "submission_result"})
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
retention_cutoff = now - timedelta(days=self.retention_days)
retained: list[dict[str, Any]] = []
for intent in intents:
completed_text = intent.get("completedAt")
if not completed_text:
retained.append(intent)
continue
try:
completed_at = datetime.fromisoformat(str(completed_text).replace("Z", "+00:00"))
except ValueError:
retained.append(intent)
continue
if completed_at.tzinfo is None:
completed_at = completed_at.replace(tzinfo=timezone.utc)
if completed_at >= retention_cutoff:
retained.append(intent)
write_jsonl(self.intents_path, retained)
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
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
write_json(destination, sanitized_market)
elif relative == "outcomes/submissions.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 not any(
marker in key.casefold()
for marker in ("url", "token", "cookie", "authorization", "configparams")
)
}
)
write_jsonl(destination, sanitized_outcomes)
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 sync(self, phase: str) -> bool:
with self._mutex:
try:
if self._workspace is None:
raise StateSyncError("state workspace is not initialized")
checksums = self._copy_snapshot()
self.generation += 1
manifest = {
"schemaVersion": STATE_SCHEMA_VERSION,
"generation": self.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 {self.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 {self.generation}".encode("utf-8"),
author=self._author,
committer=self._author,
)
current_remote_oid = self._remote_oid()
if current_remote_oid != self._expected_remote_oid:
raise StateSyncError("state branch changed remotely; refusing to overwrite a newer snapshot")
porcelain.push(
repo,
self.remote,
refspecs=f"HEAD:refs/heads/{self.branch}",
force=True,
outstream=io.BytesIO(),
errstream=io.BytesIO(),
**self._auth_kwargs(),
)
local_oid = repo.head().decode("ascii")
if self._remote_oid() != local_oid:
raise StateSyncError("state branch verification failed after push")
self._expected_remote_oid = local_oid
self.last_sync_at = manifest["updatedAt"]
self.last_error = None
self.healthy = True
self.log(f"[state-sync] generation={self.generation} phase={phase} status=ok")
return True
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)