Switch agent discovery to ModelScope
This commit is contained in:
@@ -7,9 +7,9 @@ from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
from common import utc_now, write_json
|
||||
from hf_discovery import HuggingFaceDiscovery
|
||||
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission
|
||||
from modelhub_client import ModelHubClient
|
||||
from modelhub_client import ModelHubClient, MultiModelHubClient
|
||||
from modelscope_discovery import ModelScopeDiscovery
|
||||
from outcome_tracker import OutcomeTracker
|
||||
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
||||
from template_selector import TemplateSelector
|
||||
@@ -55,7 +55,7 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
default=4,
|
||||
help="Multiplier used when auto-deriving scan limit from quota/queue capacity",
|
||||
)
|
||||
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,
|
||||
@@ -78,16 +78,18 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
parser.add_argument("--skip-outcome-sync", action="store_true", help="Skip outcome sync from ModelHub before scanning")
|
||||
parser.add_argument("--skip-history-archive", action="store_true", help="Skip historical task archive download for this run")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Plan the day 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=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
||||
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), 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("--modelscope-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)
|
||||
return parser
|
||||
|
||||
|
||||
@@ -119,6 +121,7 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar
|
||||
history_archive_path=base_args.history_archive_path,
|
||||
history_archive_limit=base_args.history_archive_limit,
|
||||
hf_base_url=base_args.hf_base_url,
|
||||
modelscope_base_url=getattr(base_args, "modelscope_base_url", base_args.hf_base_url),
|
||||
modelhub_base_url=base_args.modelhub_base_url,
|
||||
modelhub_token=base_args.modelhub_token,
|
||||
)
|
||||
@@ -130,15 +133,20 @@ def run_daily_batches(
|
||||
waves: tuple[WaveSpec, ...] = DEFAULT_WAVES,
|
||||
now=None,
|
||||
run_fn: Callable[..., dict[str, Any]] = run_submission,
|
||||
hf_discovery: HuggingFaceDiscovery | None = None,
|
||||
hf_discovery: ModelScopeDiscovery | None = None,
|
||||
modelhub_client: ModelHubClient | None = None,
|
||||
template_selector: TemplateSelector | None = None,
|
||||
outcome_tracker: OutcomeTracker | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = now or utc_now()
|
||||
hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url)
|
||||
discovery_base_url = getattr(base_args, "modelscope_base_url", None) or base_args.hf_base_url
|
||||
hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url)
|
||||
if modelhub_client is None:
|
||||
modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
|
||||
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", []) or [])
|
||||
if len(modelhub_tokens) > 1:
|
||||
modelhub_client = MultiModelHubClient(tokens=modelhub_tokens, base_url=base_args.modelhub_base_url)
|
||||
else:
|
||||
modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
|
||||
template_selector = template_selector or TemplateSelector()
|
||||
|
||||
daily_run_dir = Path(base_args.daily_runs_dir) / now.strftime("%Y%m%dT%H%M%SZ")
|
||||
@@ -284,8 +292,8 @@ def main(argv: list[str] | None = None) -> int:
|
||||
args = parser.parse_args(argv)
|
||||
ensure_tokens(args)
|
||||
log(
|
||||
f"[daily] hf_token={'set' if bool(args.hf_token) else 'missing'} "
|
||||
f"xc_token={'set' if bool(args.modelhub_token) else 'missing'}"
|
||||
f"[daily] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} "
|
||||
f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
|
||||
)
|
||||
summary = run_daily_batches(base_args=args)
|
||||
print(f"daily_run_dir={summary['dailyRunDir']}")
|
||||
|
||||
Reference in New Issue
Block a user