from __future__ import annotations import fcntl import hashlib import json import os import shutil import stat import subprocess import tempfile import threading import uuid import re from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Iterable 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", "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, 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._askpass_dir: Path | None = None self._askpass_path: 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 if self._askpass_dir is not None: shutil.rmtree(self._askpass_dir, ignore_errors=True) self._askpass_dir = None self._askpass_path = None def _git_environment(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") if self._askpass_path is None: directory = Path(tempfile.mkdtemp(prefix="modelhub-state-askpass-")) script = directory / "askpass.py" script.write_text( "#!/usr/bin/env python3\n" "import os, sys\n" "prompt = ' '.join(sys.argv[1:]).lower()\n" "key = 'MODELHUB_STATE_GIT_USERNAME' if 'username' in prompt else 'MODELHUB_STATE_GIT_PASSWORD'\n" "print(os.environ.get(key, ''))\n", encoding="utf-8", ) script.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) self._askpass_dir = directory self._askpass_path = script env = os.environ.copy() env.update( { "GIT_ASKPASS": str(self._askpass_path), "GIT_TERMINAL_PROMPT": "0", "MODELHUB_STATE_GIT_USERNAME": self.credentials["username"], "MODELHUB_STATE_GIT_PASSWORD": self.credentials["password"], } ) return env def _git(self, *args: str, cwd: Path | None = None, check: bool = True) -> subprocess.CompletedProcess[str]: try: result = subprocess.run( ["git", *args], cwd=str(cwd or self.project_root), env=self._git_environment(), text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=45, check=False, ) except FileNotFoundError as exc: raise StateSyncError( "git executable is unavailable; rebuild the service image from the repository Dockerfile" ) from exc except subprocess.TimeoutExpired as exc: raise StateSyncError("git operation timed out after 45 seconds") from exc if check and result.returncode != 0: reason = _safe_text(result.stderr.strip() or result.stdout.strip() or f"git exited {result.returncode}") raise StateSyncError(reason) return result def _remote_oid(self) -> str | None: result = self._git("ls-remote", "--heads", self.remote, self.branch) line = result.stdout.strip().splitlines() return line[0].split()[0] if line 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: self._git( "clone", "--quiet", "--single-branch", "--branch", self.branch, self.remote, str(workspace), cwd=parent, ) else: workspace.mkdir(parents=True) self._git("init", "--quiet", cwd=workspace) self._git("checkout", "--orphan", self.branch, cwd=workspace) self._git("remote", "add", "origin", self.remote, cwd=workspace) self._git("config", "user.name", self.credentials["username"], cwd=workspace) self._git("config", "user.email", self.credentials["email"], cwd=workspace) 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 commits = self._git( "rev-list", f"--max-count={self.history_depth}", "HEAD", cwd=self._workspace, ).stdout.splitlines()[1:] for commit in commits: self._git("checkout", "--quiet", "--detach", commit, cwd=self._workspace) try: manifest = self._validate_manifest(self._workspace) self.log(f"[state-recovery] fallback_commit={commit[: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}") 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) self._git("add", "--all", cwd=self._workspace) diff = self._git("diff", "--cached", "--quiet", cwd=self._workspace, check=False) if diff.returncode == 0: self.healthy = True return True self._git("commit", "--quiet", "-m", f"state: generation {self.generation} ({phase})", cwd=self._workspace) commit_count = int( self._git("rev-list", "--count", "HEAD", cwd=self._workspace).stdout.strip() or "0" ) if commit_count > self.history_depth: compact_branch = f"state-compact-{uuid.uuid4().hex[:8]}" self._git("checkout", "--quiet", "--orphan", compact_branch, cwd=self._workspace) self._git("rm", "--quiet", "--cached", "-r", ".", cwd=self._workspace, check=False) self._git("add", "--all", cwd=self._workspace) self._git( "commit", "--quiet", "-m", f"state: compacted generation {self.generation}", cwd=self._workspace, ) lease = ( f"--force-with-lease=refs/heads/{self.branch}:{self._expected_remote_oid}" if self._expected_remote_oid else f"--force-with-lease=refs/heads/{self.branch}:" ) self._git("push", "--quiet", lease, "origin", f"HEAD:refs/heads/{self.branch}", cwd=self._workspace) self._expected_remote_oid = self._git("rev-parse", "HEAD", cwd=self._workspace).stdout.strip() 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)