Rebuild agent on original pooled runner
This commit is contained in:
@@ -9,7 +9,7 @@ from typing import Any, Callable
|
||||
from common import parse_datetime, utc_now, write_json, write_jsonl
|
||||
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log
|
||||
from history_stats import is_success, update_history_archive
|
||||
from modelhub_client import ModelHubClient
|
||||
from modelhub_client import ModelHubClient, ModelHubClientPool
|
||||
from models import HFModelSummary
|
||||
from poll_runner import DEFAULT_POLL_RUNS_DIR, run_poll_loop
|
||||
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
||||
@@ -63,7 +63,7 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
action="store_true",
|
||||
help="Accepted for backward compatibility; this runner always restricts to historically successful GPUs.",
|
||||
)
|
||||
parser.add_argument("--min-downloads", type=int, default=50, help="Minimum Hugging Face download threshold")
|
||||
parser.add_argument("--min-downloads", type=int, default=50, help="Minimum ModelScope download threshold")
|
||||
parser.add_argument(
|
||||
"--history-stats-threshold",
|
||||
type=int,
|
||||
@@ -72,7 +72,7 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
)
|
||||
parser.add_argument("--read-concurrency", type=int, default=4, help="Concurrency for read-only remote calls")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Plan the queue without creating tasks")
|
||||
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing HF_TOKEN/XC_TOKEN")
|
||||
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN")
|
||||
parser.add_argument("--runs-dir", default="runs", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--ledger-path", default="ledger/submissions.jsonl", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
||||
@@ -80,10 +80,11 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--outcomes-path", default="outcomes/submissions.jsonl", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--hf-base-url", default=os.getenv("HF_BASE_URL", "https://huggingface.co"), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--hf-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--modelhub-base-url", default=os.getenv("MODELHUB_BASE_URL", "https://modelhub.org.cn"), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--modelhub-token", default=os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN"), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--hf-token", default=os.getenv("HF_TOKEN"), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--modelscope-token", default=os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN"), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--poll-interval-seconds", type=int, default=60, help="Sleep between polling cycles when no slots are available")
|
||||
parser.add_argument("--idle-interval-seconds", type=int, default=30, help="Sleep between cycles when a scan submits nothing")
|
||||
parser.add_argument(
|
||||
@@ -96,7 +97,10 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
return parser
|
||||
|
||||
|
||||
def make_modelhub_client(args: argparse.Namespace) -> ModelHubClient:
|
||||
def make_modelhub_client(args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool:
|
||||
tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token]))
|
||||
if len(tokens) > 1:
|
||||
return ModelHubClientPool([ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in tokens])
|
||||
return ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url)
|
||||
|
||||
|
||||
@@ -115,6 +119,9 @@ def model_id_from_task(task: dict[str, Any]) -> str | None:
|
||||
if value:
|
||||
return str(value).strip()
|
||||
address = str(task.get("modelAddress") or "").strip().rstrip("/")
|
||||
marker = "modelscope.cn/models/"
|
||||
if marker in address:
|
||||
return address.split(marker, 1)[1].strip("/")
|
||||
marker = "huggingface.co/"
|
||||
if marker in address:
|
||||
return address.split(marker, 1)[1].strip("/")
|
||||
@@ -243,7 +250,7 @@ def run_success_history_poll(
|
||||
*,
|
||||
now=None,
|
||||
run_fn: Callable[..., dict[str, Any]] = run_poll_loop,
|
||||
modelhub_client: ModelHubClient | None = None,
|
||||
modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
|
||||
template_selector: TemplateSelector | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = now or utc_now()
|
||||
@@ -315,8 +322,9 @@ def main(argv: list[str] | None = None) -> int:
|
||||
args = parser.parse_args(argv)
|
||||
ensure_tokens(args)
|
||||
log(
|
||||
f"[success-history] hf_token={'set' if bool(args.hf_token) else 'missing'} "
|
||||
f"xc_token={'set' if bool(args.modelhub_token) else 'missing'}"
|
||||
f"[success-history] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} "
|
||||
f"xc_token={'set' if bool(args.modelhub_token) else 'missing'} "
|
||||
f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
|
||||
)
|
||||
summary = run_success_history_poll(args)
|
||||
print(f"poll_run_dir={summary['pollRunDir']}")
|
||||
|
||||
Reference in New Issue
Block a user