Rebuild agent on original pooled runner
This commit is contained in:
@@ -9,9 +9,9 @@ from typing import Any, Callable
|
||||
|
||||
from common import utc_now, write_json
|
||||
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches
|
||||
from hf_discovery import HuggingFaceDiscovery
|
||||
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR
|
||||
from modelhub_client import ModelHubClient, MultiModelHubClient
|
||||
from modelscope_discovery import ModelScopeDiscovery
|
||||
from modelhub_client import ModelHubClient, ModelHubClientPool
|
||||
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
|
||||
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
||||
from template_selector import TemplateSelector
|
||||
@@ -64,7 +64,6 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
parser.add_argument("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS)
|
||||
parser.add_argument("--hf-base-url", default="https://modelscope.cn", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--modelscope-base-url", default="https://modelscope.cn", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--modelhub-base-url", default="https://modelhub.org.cn", help=argparse.SUPPRESS)
|
||||
parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS)
|
||||
parser.add_argument("--hf-token", default=None, help=argparse.SUPPRESS)
|
||||
@@ -86,10 +85,11 @@ def _make_cycle_args(base_args: argparse.Namespace) -> argparse.Namespace:
|
||||
return cycle_args
|
||||
|
||||
|
||||
def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient:
|
||||
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", []) or [])
|
||||
def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool:
|
||||
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
|
||||
if len(modelhub_tokens) > 1:
|
||||
return MultiModelHubClient(tokens=modelhub_tokens, base_url=base_args.modelhub_base_url)
|
||||
clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in modelhub_tokens]
|
||||
return ModelHubClientPool(clients)
|
||||
return ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
|
||||
|
||||
|
||||
@@ -98,14 +98,13 @@ def run_poll_loop(
|
||||
base_args: argparse.Namespace,
|
||||
now=None,
|
||||
run_fn: Callable[..., dict[str, Any]] = run_daily_batches,
|
||||
hf_discovery: ModelScopeDiscovery | None = None,
|
||||
modelhub_client: ModelHubClient | None = None,
|
||||
hf_discovery: HuggingFaceDiscovery | None = None,
|
||||
modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
|
||||
template_selector: TemplateSelector | None = None,
|
||||
outcome_tracker: OutcomeTracker | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = now or utc_now()
|
||||
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)
|
||||
hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url)
|
||||
modelhub_client = modelhub_client or _build_modelhub_client(base_args)
|
||||
template_selector = template_selector or TemplateSelector()
|
||||
|
||||
@@ -251,6 +250,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
|
||||
log(
|
||||
f"[poll] 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_poll_loop(base_args=args)
|
||||
|
||||
Reference in New Issue
Block a user