Rebuild agent on original pooled runner
This commit is contained in:
CoolBoy
@@ -8,13 +8,20 @@ from defaults import EMBEDDED_HF_TOKEN, EMBEDDED_MODELHUB_XC_TOKENS, EMBEDDED_MO
|
||||
|
||||
|
||||
DEFAULT_KEY_PATH = Path("KEY.md")
|
||||
DEFAULT_KEYS_PATH = Path("KEYS.md")
|
||||
MODULE_DIR = Path(__file__).resolve().parent
|
||||
XC_TOKEN_ENV_NAMES = ("MODELHUB_XC_TOKEN", "XC_TOKEN", "MODELHUB_TOKEN")
|
||||
XC_TOKEN_LIST_ENV_NAMES = ("MODELHUB_XC_TOKENS", "XC_TOKENS", "MODELHUB_TOKENS")
|
||||
JWT_TOKEN_ENV_NAMES = ("MODELHUB_JWT_TOKEN", "JWT_TOKEN")
|
||||
MODELSCOPE_TOKEN_ENV_NAMES = ("MODELSCOPE_API_TOKEN", "MODELSCOPE_TOKEN")
|
||||
|
||||
|
||||
def _token_sort_key(key: str) -> tuple[int, str]:
|
||||
suffix = key.removeprefix("XC_TOKEN")
|
||||
if not suffix:
|
||||
return (0, key)
|
||||
if suffix.isdigit():
|
||||
return (int(suffix), key)
|
||||
return (10_000, key)
|
||||
|
||||
|
||||
def load_key_file(path: Path) -> dict[str, str]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
@@ -28,109 +35,108 @@ def load_key_file(path: Path) -> dict[str, str]:
|
||||
return loaded
|
||||
|
||||
|
||||
def first_value(values: dict[str, str], names: tuple[str, ...]) -> str | None:
|
||||
for name in names:
|
||||
value = os.getenv(name) or values.get(name)
|
||||
if value:
|
||||
return value.strip()
|
||||
return None
|
||||
def load_key_files(*paths: Path) -> dict[str, str]:
|
||||
loaded: dict[str, str] = {}
|
||||
for path in paths:
|
||||
if path.exists():
|
||||
loaded.update(load_key_file(path))
|
||||
return loaded
|
||||
|
||||
|
||||
def split_token_list(value: str | None) -> list[str]:
|
||||
def load_modelhub_tokens(values: dict[str, str]) -> list[str]:
|
||||
tokens: list[tuple[str, str]] = []
|
||||
seen: set[str] = set()
|
||||
for key, value in values.items():
|
||||
if key == "MODELHUB_XC_TOKEN" or key.startswith("XC_TOKEN"):
|
||||
token = value.strip()
|
||||
if token and token not in seen:
|
||||
seen.add(token)
|
||||
tokens.append((key, token))
|
||||
tokens.sort(key=lambda item: _token_sort_key(item[0]))
|
||||
return [token for _, token in tokens]
|
||||
|
||||
|
||||
def _split_token_list(value: str | None) -> list[str]:
|
||||
if not value:
|
||||
return []
|
||||
tokens: list[str] = []
|
||||
for normalized in value.replace(",", "\n").replace(";", "\n").splitlines():
|
||||
token = normalized.strip()
|
||||
for raw in value.replace(",", "\n").replace(";", "\n").splitlines():
|
||||
token = raw.strip()
|
||||
if token:
|
||||
tokens.append(token)
|
||||
return tokens
|
||||
|
||||
|
||||
def discover_key_paths(primary_path: Path) -> list[Path]:
|
||||
candidates: list[Path] = []
|
||||
search_dirs = [primary_path.parent, primary_path.parent.parent]
|
||||
for path in [primary_path, primary_path.parent / "KEYS.md"]:
|
||||
if path not in candidates:
|
||||
candidates.append(path)
|
||||
for directory in search_dirs:
|
||||
if not directory.exists():
|
||||
continue
|
||||
for path in sorted(directory.glob("KEYS*")):
|
||||
if path.is_file() and path not in candidates:
|
||||
candidates.append(path)
|
||||
return candidates
|
||||
|
||||
|
||||
def collect_modelhub_tokens(values_by_path: list[dict[str, str]], explicit_token: str | None) -> list[str]:
|
||||
tokens: list[str] = []
|
||||
|
||||
def add(value: str | None) -> None:
|
||||
for token in split_token_list(value):
|
||||
if token and token not in tokens:
|
||||
tokens.append(token)
|
||||
|
||||
add(explicit_token)
|
||||
for env_name in XC_TOKEN_LIST_ENV_NAMES:
|
||||
add(os.getenv(env_name))
|
||||
for env_name in XC_TOKEN_ENV_NAMES:
|
||||
add(os.getenv(env_name))
|
||||
|
||||
for values in values_by_path:
|
||||
for key, value in values.items():
|
||||
normalized = key.strip().upper()
|
||||
if normalized in XC_TOKEN_ENV_NAMES or normalized.startswith("XC_TOKEN") or normalized.startswith("MODELHUB_XC_TOKEN"):
|
||||
add(value)
|
||||
|
||||
for token in EMBEDDED_MODELHUB_XC_TOKENS:
|
||||
add(token)
|
||||
return tokens
|
||||
def _add_token(tokens: list[str], token: str | None) -> None:
|
||||
for item in _split_token_list(token):
|
||||
if item and item not in tokens:
|
||||
tokens.append(item)
|
||||
|
||||
|
||||
def ensure_tokens(args: argparse.Namespace) -> None:
|
||||
primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH))
|
||||
supplemental_key_path = primary_key_path.with_name(DEFAULT_KEYS_PATH.name)
|
||||
if not primary_key_path.exists():
|
||||
fallback_paths = [
|
||||
for candidate in (
|
||||
MODULE_DIR / primary_key_path.name,
|
||||
primary_key_path.parent.parent / primary_key_path.name,
|
||||
MODULE_DIR.parent / primary_key_path.name,
|
||||
]
|
||||
for fallback_path in fallback_paths:
|
||||
if fallback_path.exists():
|
||||
primary_key_path = fallback_path
|
||||
):
|
||||
if candidate.exists():
|
||||
primary_key_path = candidate
|
||||
supplemental_key_path = candidate.with_name(DEFAULT_KEYS_PATH.name)
|
||||
break
|
||||
values_by_path = [load_key_file(path) for path in discover_key_paths(primary_key_path) if path.exists()]
|
||||
values: dict[str, str] = {}
|
||||
for loaded in values_by_path:
|
||||
values.update(loaded)
|
||||
values = load_key_files(primary_key_path, supplemental_key_path)
|
||||
|
||||
if not getattr(args, "hf_token", None):
|
||||
args.hf_token = os.getenv("HF_TOKEN") or values.get("HF_TOKEN")
|
||||
if not getattr(args, "hf_token", None):
|
||||
args.hf_token = EMBEDDED_HF_TOKEN
|
||||
args.hf_token = values.get("HF_TOKEN") or EMBEDDED_HF_TOKEN
|
||||
if not getattr(args, "modelscope_token", None):
|
||||
args.modelscope_token = first_value(values, MODELSCOPE_TOKEN_ENV_NAMES)
|
||||
if not getattr(args, "modelscope_token", None):
|
||||
args.modelscope_token = EMBEDDED_MODELSCOPE_TOKEN
|
||||
args.modelscope_token = (
|
||||
os.getenv("MODELSCOPE_API_TOKEN")
|
||||
or os.getenv("MODELSCOPE_TOKEN")
|
||||
or values.get("MODELSCOPE_API_TOKEN")
|
||||
or values.get("MODELSCOPE_TOKEN")
|
||||
or EMBEDDED_MODELSCOPE_TOKEN
|
||||
)
|
||||
|
||||
modelhub_tokens = collect_modelhub_tokens(values_by_path, getattr(args, "modelhub_token", None))
|
||||
modelhub_token = modelhub_tokens[0] if modelhub_tokens else None
|
||||
jwt_token = first_value(values, JWT_TOKEN_ENV_NAMES)
|
||||
tokens = load_modelhub_tokens(values)
|
||||
env_token_items: list[tuple[str, str]] = []
|
||||
for key, value in os.environ.items():
|
||||
if key == "MODELHUB_XC_TOKEN" or key.startswith("XC_TOKEN"):
|
||||
token = value.strip()
|
||||
if token:
|
||||
env_token_items.append((key, token))
|
||||
env_token_items.sort(key=lambda item: _token_sort_key(item[0]))
|
||||
env_tokens = [token for _, token in env_token_items]
|
||||
|
||||
args.modelhub_token = modelhub_token
|
||||
args.modelhub_tokens = modelhub_tokens
|
||||
modelhub_token = getattr(args, "modelhub_token", None)
|
||||
if modelhub_token and modelhub_token not in env_tokens:
|
||||
env_tokens.insert(0, modelhub_token)
|
||||
if not env_tokens:
|
||||
env_tokens = tokens
|
||||
else:
|
||||
for token in tokens:
|
||||
if token not in env_tokens:
|
||||
env_tokens.append(token)
|
||||
for env_name in ("MODELHUB_XC_TOKENS", "XC_TOKENS", "MODELHUB_TOKENS"):
|
||||
for token in _split_token_list(os.getenv(env_name)):
|
||||
_add_token(env_tokens, token)
|
||||
for token in EMBEDDED_MODELHUB_XC_TOKENS:
|
||||
_add_token(env_tokens, token)
|
||||
|
||||
args.modelhub_tokens = env_tokens
|
||||
args.modelhub_token = env_tokens[0] if env_tokens else None
|
||||
|
||||
if args.hf_token:
|
||||
os.environ["HF_TOKEN"] = args.hf_token
|
||||
if args.modelscope_token:
|
||||
os.environ["MODELSCOPE_API_TOKEN"] = args.modelscope_token
|
||||
os.environ["MODELSCOPE_TOKEN"] = args.modelscope_token
|
||||
for index, token in enumerate(args.modelhub_tokens, start=1):
|
||||
env_name = "XC_TOKEN" if index == 1 else f"XC_TOKEN{index}"
|
||||
os.environ[env_name] = token
|
||||
if args.modelhub_token:
|
||||
os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token
|
||||
os.environ["XC_TOKEN"] = args.modelhub_token
|
||||
if args.modelhub_tokens:
|
||||
os.environ["MODELHUB_XC_TOKENS"] = ",".join(args.modelhub_tokens)
|
||||
if jwt_token:
|
||||
os.environ["MODELHUB_JWT_TOKEN"] = jwt_token
|
||||
if not args.modelhub_token and not jwt_token:
|
||||
raise ValueError("MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN is required")
|
||||
if not args.modelhub_token:
|
||||
raise ValueError("XC_TOKEN is required, either via environment or KEY.md")
|
||||
|
||||
Reference in New Issue
Block a user