From 1336a38ee4935d577617ee9bd2cfb4ef692996a1 Mon Sep 17 00:00:00 2001 From: angosr Date: Mon, 18 May 2026 14:23:18 +0000 Subject: [PATCH 1/3] scheduler: mirror champions to official hf repo --- affine/src/scheduler/champion_mirror.py | 138 ++++++++++++++++++++++++ affine/src/scheduler/flow.py | 18 +++- 2 files changed, 154 insertions(+), 2 deletions(-) create mode 100644 affine/src/scheduler/champion_mirror.py diff --git a/affine/src/scheduler/champion_mirror.py b/affine/src/scheduler/champion_mirror.py new file mode 100644 index 00000000..32e77d28 --- /dev/null +++ b/affine/src/scheduler/champion_mirror.py @@ -0,0 +1,138 @@ +"""Mirror champion checkpoints into an official HuggingFace repo.""" + +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_REPO_ENV = "AFFINE_OFFICIAL_HF_REPO" +OFFICIAL_HF_REPO_DEFAULT = "AffineFoundation/affine-champion" + + +class ChampionMirror: + """Pushes the champion's exact git commit to the official HF repo.""" + + def __init__( + self, + *, + token: Optional[str] = None, + repo_id: Optional[str] = None, + ): + self.token = (token if token is not None else os.getenv(OFFICIAL_HF_TOKEN_ENV, "")).strip() + self.repo_id = ( + repo_id + if repo_id is not None + else os.getenv(OFFICIAL_HF_REPO_ENV, OFFICIAL_HF_REPO_DEFAULT) + ).strip() + + def enabled(self) -> bool: + return bool(self.token and self.repo_id) + + async def ensure_mirrored(self, champion: ChampionRecord) -> ChampionRecord: + if not self.enabled(): + return champion + if champion.model == self.repo_id: + return champion + try: + return await asyncio.to_thread(self._mirror_sync, champion) + except Exception as e: + logger.warning( + f"ChampionMirror: mirror failed uid={champion.uid} " + f"{champion.model}@{champion.revision[:8]} -> {self.repo_id}: " + f"{type(e).__name__}: {e}" + ) + return champion + + def _mirror_sync(self, champion: ChampionRecord) -> ChampionRecord: + api = HfApi(token=self.token) + api.create_repo( + repo_id=self.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(self.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"{self.repo_id}@{champion.revision[:8]}" + ) + return replace(champion, model=self.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 _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. From f904e5eff72178ad8091c802e391cc214997d875 Mon Sep 17 00:00:00 2001 From: angosr Date: Mon, 18 May 2026 14:30:04 +0000 Subject: [PATCH 2/3] scheduler: mirror champions into affine-prefixed repos --- affine/src/scheduler/champion_mirror.py | 46 +++++++++++++++---------- 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/affine/src/scheduler/champion_mirror.py b/affine/src/scheduler/champion_mirror.py index 32e77d28..eec577d0 100644 --- a/affine/src/scheduler/champion_mirror.py +++ b/affine/src/scheduler/champion_mirror.py @@ -1,4 +1,4 @@ -"""Mirror champion checkpoints into an official HuggingFace repo.""" +"""Mirror champion checkpoints into official HuggingFace repos.""" from __future__ import annotations @@ -18,48 +18,51 @@ OFFICIAL_HF_TOKEN_ENV = "AFFINE_OFFICIAL_HF_TOKEN" -OFFICIAL_HF_REPO_ENV = "AFFINE_OFFICIAL_HF_REPO" -OFFICIAL_HF_REPO_DEFAULT = "AffineFoundation/affine-champion" +OFFICIAL_HF_NAMESPACE_ENV = "AFFINE_OFFICIAL_HF_NAMESPACE" +OFFICIAL_HF_NAMESPACE_DEFAULT = "AffineFoundation" class ChampionMirror: - """Pushes the champion's exact git commit to the official HF repo.""" + """Pushes the champion's exact git commit to an official affine repo.""" def __init__( self, *, token: Optional[str] = None, - repo_id: Optional[str] = None, + namespace: Optional[str] = None, ): self.token = (token if token is not None else os.getenv(OFFICIAL_HF_TOKEN_ENV, "")).strip() - self.repo_id = ( - repo_id - if repo_id is not None - else os.getenv(OFFICIAL_HF_REPO_ENV, OFFICIAL_HF_REPO_DEFAULT) + 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.repo_id) + return bool(self.token and self.namespace) async def ensure_mirrored(self, champion: ChampionRecord) -> ChampionRecord: if not self.enabled(): return champion - if champion.model == self.repo_id: + repo_id = _target_repo_id(champion.model, self.namespace) + if champion.model == repo_id: return champion try: - return await asyncio.to_thread(self._mirror_sync, champion) + 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]} -> {self.repo_id}: " + f"{champion.model}@{champion.revision[:8]} -> {repo_id}: " f"{type(e).__name__}: {e}" ) return champion - def _mirror_sync(self, champion: ChampionRecord) -> ChampionRecord: + def _mirror_sync( + self, champion: ChampionRecord, target_repo_id: str + ) -> ChampionRecord: api = HfApi(token=self.token) api.create_repo( - repo_id=self.repo_id, + repo_id=target_repo_id, repo_type="model", token=self.token, exist_ok=True, @@ -67,7 +70,7 @@ def _mirror_sync(self, champion: ChampionRecord) -> ChampionRecord: 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(self.repo_id, self.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) @@ -88,9 +91,9 @@ def _mirror_sync(self, champion: ChampionRecord) -> ChampionRecord: logger.info( f"ChampionMirror: mirrored champion uid={champion.uid} " f"{champion.model}@{champion.revision[:8]} -> " - f"{self.repo_id}@{champion.revision[:8]}" + f"{target_repo_id}@{champion.revision[:8]}" ) - return replace(champion, model=self.repo_id) + return replace(champion, model=target_repo_id) def _run_git( self, @@ -122,6 +125,13 @@ def _hf_git_url(repo_id: str, token: Optional[str]) -> str: return f"https://huggingface.co/{repo_id}" +def _target_repo_id(source_repo_id: str, namespace: str) -> str: + name = source_repo_id.rstrip("/").split("/")[-1] + if not name.startswith("affine-"): + 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 "-" From 9029eb054cc1268c6f2fb51319173cc9299db66d Mon Sep 17 00:00:00 2001 From: angosr Date: Mon, 18 May 2026 14:31:27 +0000 Subject: [PATCH 3/3] scheduler: use uppercase hotkey for champion mirror repo --- affine/src/scheduler/champion_mirror.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/affine/src/scheduler/champion_mirror.py b/affine/src/scheduler/champion_mirror.py index eec577d0..d679edb2 100644 --- a/affine/src/scheduler/champion_mirror.py +++ b/affine/src/scheduler/champion_mirror.py @@ -44,7 +44,7 @@ def enabled(self) -> bool: async def ensure_mirrored(self, champion: ChampionRecord) -> ChampionRecord: if not self.enabled(): return champion - repo_id = _target_repo_id(champion.model, self.namespace) + repo_id = _target_repo_id(champion.hotkey, self.namespace) if champion.model == repo_id: return champion try: @@ -125,10 +125,12 @@ def _hf_git_url(repo_id: str, token: Optional[str]) -> str: return f"https://huggingface.co/{repo_id}" -def _target_repo_id(source_repo_id: str, namespace: str) -> str: - name = source_repo_id.rstrip("/").split("/")[-1] - if not name.startswith("affine-"): - name = f"affine-{name}" +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}"