diff --git a/affine/src/scheduler/champion_mirror.py b/affine/src/scheduler/champion_mirror.py new file mode 100644 index 00000000..d679edb2 --- /dev/null +++ b/affine/src/scheduler/champion_mirror.py @@ -0,0 +1,150 @@ +"""Mirror champion checkpoints into official HuggingFace repos.""" + +from __future__ import annotations + +import asyncio +import os +import re +import subprocess +import tempfile +from dataclasses import replace +from typing import Optional +from urllib.parse import quote + +from huggingface_hub import HfApi + +from affine.core.setup import logger +from affine.src.scorer.window_state import ChampionRecord + + +OFFICIAL_HF_TOKEN_ENV = "AFFINE_OFFICIAL_HF_TOKEN" +OFFICIAL_HF_NAMESPACE_ENV = "AFFINE_OFFICIAL_HF_NAMESPACE" +OFFICIAL_HF_NAMESPACE_DEFAULT = "AffineFoundation" + + +class ChampionMirror: + """Pushes the champion's exact git commit to an official affine repo.""" + + def __init__( + self, + *, + token: Optional[str] = None, + namespace: Optional[str] = None, + ): + self.token = (token if token is not None else os.getenv(OFFICIAL_HF_TOKEN_ENV, "")).strip() + self.namespace = ( + namespace + if namespace is not None + else os.getenv(OFFICIAL_HF_NAMESPACE_ENV, OFFICIAL_HF_NAMESPACE_DEFAULT) + ).strip() + + def enabled(self) -> bool: + return bool(self.token and self.namespace) + + async def ensure_mirrored(self, champion: ChampionRecord) -> ChampionRecord: + if not self.enabled(): + return champion + repo_id = _target_repo_id(champion.hotkey, self.namespace) + if champion.model == repo_id: + return champion + try: + return await asyncio.to_thread(self._mirror_sync, champion, repo_id) + except Exception as e: + logger.warning( + f"ChampionMirror: mirror failed uid={champion.uid} " + f"{champion.model}@{champion.revision[:8]} -> {repo_id}: " + f"{type(e).__name__}: {e}" + ) + return champion + + def _mirror_sync( + self, champion: ChampionRecord, target_repo_id: str + ) -> ChampionRecord: + api = HfApi(token=self.token) + api.create_repo( + repo_id=target_repo_id, + repo_type="model", + token=self.token, + exist_ok=True, + ) + with tempfile.TemporaryDirectory(prefix="affine-champion-") as tmp: + repo_dir = os.path.join(tmp, "repo") + source_url = _hf_git_url(champion.model, os.getenv("HF_TOKEN")) + target_url = _hf_git_url(target_repo_id, self.token) + self._run_git(["clone", source_url, repo_dir], cwd=tmp) + self._run_git(["checkout", "--detach", champion.revision], cwd=repo_dir) + self._run_git(["lfs", "install", "--local"], cwd=repo_dir) + self._run_git(["lfs", "fetch", "origin", champion.revision], cwd=repo_dir) + self._run_git(["lfs", "checkout"], cwd=repo_dir) + self._run_git(["remote", "add", "official", target_url], cwd=repo_dir) + branch = _branch_name(champion) + self._run_git( + ["push", "official", f"HEAD:refs/heads/{branch}"], + cwd=repo_dir, + timeout=3600, + ) + self._run_git( + ["push", "--force", "official", "HEAD:refs/heads/main"], + cwd=repo_dir, + timeout=3600, + ) + logger.info( + f"ChampionMirror: mirrored champion uid={champion.uid} " + f"{champion.model}@{champion.revision[:8]} -> " + f"{target_repo_id}@{champion.revision[:8]}" + ) + return replace(champion, model=target_repo_id) + + def _run_git( + self, + args: list[str], + *, + cwd: str, + timeout: int = 1800, + ) -> None: + cmd = ["git", *args] + result = subprocess.run( + cmd, + cwd=cwd, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=timeout, + check=False, + ) + if result.returncode != 0: + raise RuntimeError( + f"git {_mask_args(args)} failed with code {result.returncode}: " + f"{_mask_text(result.stderr or result.stdout)}" + ) + + +def _hf_git_url(repo_id: str, token: Optional[str]) -> str: + if token: + return f"https://__token__:{quote(token, safe='')}@huggingface.co/{repo_id}" + return f"https://huggingface.co/{repo_id}" + + +def _target_repo_id(hotkey: str, namespace: str) -> str: + name = "".join( + c if c.isalnum() or c in {"-", "_"} else "-" + for c in hotkey.strip() + ).strip("-_").upper() + name = f"affine-{name}" + return f"{namespace}/{name}" + + +def _branch_name(champion: ChampionRecord) -> str: + safe_hotkey = "".join( + c if c.isalnum() or c in {"-", "_"} else "-" + for c in champion.hotkey[:16] + ).strip("-_") or "hotkey" + return f"champions/uid-{champion.uid}-{safe_hotkey}-{champion.revision[:12]}" + + +def _mask_args(args: list[str]) -> str: + return " ".join(_mask_text(arg) for arg in args) + + +def _mask_text(text: str) -> str: + return re.sub(r"__token__:[^@\s]+", "__token__:***", text) diff --git a/affine/src/scheduler/flow.py b/affine/src/scheduler/flow.py index b00535b1..e7b9994a 100644 --- a/affine/src/scheduler/flow.py +++ b/affine/src/scheduler/flow.py @@ -65,6 +65,7 @@ ) from . import targon as targon_lifecycle +from .champion_mirror import ChampionMirror # ---- deploy_fn contract: signaling exception ------------------------------- @@ -211,6 +212,7 @@ def __init__( self._list_valid_miners = list_valid_miners_fn self._list_current_miners = list_current_miners_fn or list_valid_miners_fn self._list_active_endpoint_names = list_active_endpoint_names_fn + self._champion_mirror = ChampionMirror() # ---- entry point ------------------------------------------------------ @@ -400,7 +402,7 @@ async def _read_synced_champion(self) -> Optional[ChampionRecord]: f"FlowScheduler: champion hotkey={champ.hotkey[:10]} " f"uid synced {old_uid} -> {current_uid}" ) - return champ + return await self._ensure_champion_mirrored(champ) if champ.uid != -1: old_uid = champ.uid @@ -412,7 +414,15 @@ async def _read_synced_champion(self) -> Optional[ChampionRecord]: f"and burning weight until it re-registers " f"(old_uid={old_uid})" ) - return champ + return await self._ensure_champion_mirrored(champ) + + async def _ensure_champion_mirrored( + self, champ: ChampionRecord, + ) -> ChampionRecord: + mirrored = await self._champion_mirror.ensure_mirrored(champ) + if mirrored.model != champ.model or mirrored.revision != champ.revision: + await self.state.set_champion(mirrored) + return mirrored async def _refresh_task_ids( self, current_block: int, envs: Mapping[str, EnvConfig], @@ -502,6 +512,7 @@ async def _cold_start(self, current_block: int) -> None: revision=candidate.revision, model=candidate.model, since_block=current_block, ) + new_champ = await self._champion_mirror.ensure_mirrored(new_champ) await self.state.set_champion(new_champ) await self.queue.mark_terminated( candidate.uid, @@ -1331,6 +1342,9 @@ async def _decide( deployments=list(battle.deployments), since_block=current_block, ) + new_champion = await self._champion_mirror.ensure_mirrored( + new_champion + ) # Order: write the canonical champion record FIRST so any # concurrent reader sees the new identity before either miners # row flips.