submmit/modelhub_submmit_api/state_sync.py

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)