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 "" 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)