overlay/scripts/hf_routing.py

from __future__ import annotations

import os
from dataclasses import dataclass

from huggingface_hub import HfApi


_OWNER_ALIASES = {
    'jack': 'jackoatmon',
    'jackoatmon': 'jackoatmon',
    'icarus': 'icarus112',
    'icarus112': 'icarus112',
}


def _normalize_owner(value: str | None) -> str | None:
    if not value:
        return None
    normalized = value.strip().lower().lstrip('@')
    if not normalized:
        return None
    return _OWNER_ALIASES.get(normalized, normalized)


def _owner_from_env() -> str | None:
    for key in ('FEATHER_HF_OWNER', 'FEATHER_HF_NAMESPACE_OWNER', 'FEATHER_HF_PROFILE'):
        owner = _normalize_owner(os.environ.get(key))
        if owner:
            return owner
    return None


def resolve_owner(token: str | None = None) -> str:
    """Resolve active HF owner in a collaborator-safe way.

    Resolution precedence:
      1) explicit env owner override (FEATHER_HF_OWNER/...)
      2) Hugging Face `whoami` from HF_TOKEN (unless disabled)
      3) default to jackoatmon
    """
    owner = _owner_from_env()
    if owner:
        return owner

    if os.environ.get('FEATHER_HF_DISABLE_WHOAMI', '0') != '1':
        active_token = token or os.environ.get('HF_TOKEN')
        if active_token:
            try:
                info = HfApi(token=active_token).whoami(token=active_token)
                if isinstance(info, dict):
                    whoami_owner = _normalize_owner(info.get('name'))
                    if whoami_owner:
                        return whoami_owner
            except Exception:
                # Fail open to deterministic defaults for offline/dry-run tests.
                pass

    return 'jackoatmon'


@dataclass(frozen=True)
class HfRouting:
    owner: str
    space_repo: str
    output_repo: str
    retina_cache_repo: str
    job_namespace: str


def resolve_routing(token: str | None = None) -> HfRouting:
    owner = resolve_owner(token=token)

    space_name = os.environ.get('FEATHER_HF_SPACE_NAME', 'feather-runtime')
    output_name = os.environ.get('FEATHER_HF_OUTPUT_REPO_NAME', 'feather-pretrain-checkpoints')
    retina_name = os.environ.get('FEATHER_HF_RETINA_REPO_NAME', 'feather-retina-cache')

    space_repo = os.environ.get('FEATHER_HF_SPACE_REPO') or f'{owner}/{space_name}'
    output_repo = os.environ.get('FEATHER_HF_OUTPUT_REPO') or f'{owner}/{output_name}'
    retina_cache_repo = os.environ.get('FEATHER_HF_RETINA_CACHE_REPO') or f'{owner}/{retina_name}'
    job_namespace = os.environ.get('FEATHER_HF_JOB_NAMESPACE') or owner

    return HfRouting(
        owner=owner,
        space_repo=space_repo,
        output_repo=output_repo,
        retina_cache_repo=retina_cache_repo,
        job_namespace=job_namespace,
    )