"""HuggingFace Token 共享 helper - 提供 HF 账户切换和 token 获取 提供以下公共 API: - get_hf_token_interactive() 必要时 prompt 后获取 token - get_hf_username() 获取当前 HF 用户名 设计: - 与 bootstrap-hf.sh 切换流程保持一致 - 非 TTY 时 (stdin.isatty() == False) 不进入 prompt 循环，避免死锁 - login 成功后调用 hf auth whoami 验证，切换时输出 switched 信息 - 不实现 RESTORE_HF_LOGIN_* 恢复机制（单命令工具不需要） 用法: import sys, os sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from _hf_token import get_hf_token_interactive token = get_hf_token_interactive() if not token: sys.exit(1) """ from __future__ import annotations import getpass import os import re import subprocess import sys from typing import Optional _HF_CLI_TIMEOUT = 10 # hf CLI 调用超时（秒） def _parse_whoami_username(output: str) -> str: """解析 hf auth whoami 输出，提取 username。 接受两种格式: "user: " "Logged in as " """ match = re.search(r"user:\s*(\S+)", output or "") if not match: match = re.search(r"[Ll]ogged in as\s+(\S+)", output or "") return match.group(1) if match else "" _HF_EXTRA_PATH = "/root/.local/bin" def _run_hf(args: list[str]) -> tuple[int, str, str]: """运行 hf CLI 子命令并返回 (returncode, stdout, stderr)。""" import shutil as _shutil # Check if hf is available at all before trying subprocess if not _shutil.which("hf") and not _shutil.which(os.path.join(_HF_EXTRA_PATH, "hf")): return (1, "", "hf not found in PATH") try: result = subprocess.run( ["hf", *args], capture_output=True, text=True, timeout=_HF_CLI_TIMEOUT, ) return result.returncode, result.stdout, result.stderr except FileNotFoundError: pass except (subprocess.TimeoutExpired) as exc: return 1, "", str(exc) # hf not found — retry with /root/.local/bin in PATH (handles cron/systemd env) extra_env = os.environ.copy() extra_env["PATH"] = _HF_EXTRA_PATH + os.pathsep + extra_env.get("PATH", "") try: result = subprocess.run( ["hf", *args], capture_output=True, text=True, timeout=_HF_CLI_TIMEOUT, env=extra_env, ) return result.returncode, result.stdout, result.stderr except (subprocess.TimeoutExpired, FileNotFoundError) as exc: return 1, "", str(exc) def _read_cached_token() -> str: """从 env 或 cache 文件读取 token。 优先级: HUGGINGFACE_HUB_TOKEN > HF_TOKEN > HF_TOKEN_FILE > ~/.cache/huggingface/token """ for var in ("HUGGINGFACE_HUB_TOKEN", "HF_TOKEN"): val = os.environ.get(var, "").strip() if val: return val token_file = os.environ.get("HF_TOKEN_FILE") if not token_file: home = os.environ.get("HOME") or "/root" token_file = os.path.join(home, ".cache", "huggingface", "token") if os.path.exists(token_file): try: with open(token_file) as fh: return fh.read().strip() except OSError: return "" return "" def get_hf_username() -> Optional[str]: """获取当前 HF 用户名。失败返回 None。""" rc, stdout, _ = _run_hf(["auth", "whoami"]) if rc != 0: return None username = _parse_whoami_username(stdout) if username: return username # Python fallback: 用 cache token 调 HfApi token = _read_cached_token() if not token: return None try: from huggingface_hub import HfApi api = HfApi(token=token) data = api.whoami(token=token) name = data.get("name", "") if isinstance(data, dict) else "" return (name or "").strip() or None except Exception: return None def _prompt_for_token() -> Optional[str]: """交互式读取一个 token。已确认 stdin 是 TTY。""" try: token = getpass.getpass("Enter HF_TOKEN: ").strip() except (EOFError, KeyboardInterrupt): print("\n[WARN] Token input cancelled", file=sys.stderr) return None if not token: print("[WARN] Empty token, please try again", file=sys.stderr) return None return token def get_hf_token_interactive() -> Optional[str]: """获取 HF token，支持账户切换。 流程: 1. 若已登录（hf auth whoami 成功）: - TTY: 询问 "Use this user? [Y/n]"，如拒绝则 prompt 新 token 并切换 - 非 TTY: 静默使用当前用户 2. 未登录: - 尝试 env / cache 文件 - 仍无: TTY 下 prompt 新 token 登录；非 TTY 返回 None 3. login 成功后调用 hf auth whoami 验证，输出 switched 信息 成功返回 token 字符串；非交互且无 token 时返回 None。 """ current_user: Optional[str] = None is_tty = sys.stdin.isatty() # ---- 步骤 1: 检测是否已登录 ---- rc, stdout, _ = _run_hf(["auth", "whoami"]) if rc == 0: current_user = _parse_whoami_username(stdout) or None if current_user: print(f"[INFO] HF CLI is logged in as: {current_user}") if is_tty: response = input("Use this user? (Y/n) [Y]: ").strip().lower() if response in ("", "y", "yes"): token = _read_cached_token() if token: return token # cache 缺失却 whoami 成功，尝试从环境拿到 env_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get( "HF_TOKEN", "" ) if env_token.strip(): return env_token.strip() # 用户选择切换 → 进入步骤 3 else: # 非 TTY：静默使用当前 token = _read_cached_token() if token: return token env_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get( "HF_TOKEN", "" ) if env_token.strip(): return env_token.strip() # ---- 步骤 2: 未登录时尝试 env / cache ---- if not current_user: token = _read_cached_token() if token: return token # ---- 步骤 3: prompt 新 token 并登录 ---- if not is_tty: print( "[ERROR] Cannot prompt for HF_TOKEN in non-interactive mode. " "Set HF_TOKEN env var or run: hf auth login", file=sys.stderr, ) return None print("[INFO] HF CLI is not logged in. Please provide a HF_TOKEN.") while True: token = _prompt_for_token() if not token: continue login_rc, _, login_stderr = _run_hf(["auth", "login", "--token", token]) if login_rc != 0: print( f"[ERROR] HF login failed: {login_stderr.strip() or 'unknown error'}", file=sys.stderr, ) continue # login 成功 → whoami 验证 verify_rc, verify_stdout, _ = _run_hf(["auth", "whoami"]) if verify_rc != 0: print( "[ERROR] HF login succeeded but whoami verification failed", file=sys.stderr, ) continue new_user = _parse_whoami_username(verify_stdout) or "unknown" if current_user and new_user != current_user: print(f"[INFO] HF account switched: {current_user} -> {new_user}") elif current_user: print(f"[INFO] HF login successful as: {new_user}") else: print(f"[INFO] HF login successful as: {new_user}") return token if __name__ == "__main__": # 简单自测入口 token = get_hf_token_interactive() if token: print(f"Got token (len={len(token)})") sys.exit(0) sys.exit(1)