| """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 |
|
|
|
|
| def _parse_whoami_username(output: str) -> str: |
| """解析 hf auth whoami 输出，提取 username。 |
| |
| 接受两种格式: |
| "user: <name>" |
| "Logged in as <name>" |
| """ |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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 |
| |
| 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() |
|
|
| |
| 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 |
| |
| env_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get( |
| "HF_TOKEN", "" |
| ) |
| if env_token.strip(): |
| return env_token.strip() |
| |
| else: |
| |
| 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() |
|
|
| |
| if not current_user: |
| token = _read_cached_token() |
| if token: |
| return 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 |
|
|
| |
| 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) |
|
|
