650 lines
30 KiB
Python
650 lines
30 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
from common import read_jsonl, utc_now, write_json
|
|
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches
|
|
from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH
|
|
from hf_discovery import HuggingFaceDiscovery
|
|
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir
|
|
from market_intelligence import (
|
|
DEFAULT_FETCH_WORKERS,
|
|
DEFAULT_FRAMEWORK_MIN_SAMPLES,
|
|
DEFAULT_FRAMEWORK_REFRESH_SECONDS,
|
|
DEFAULT_MARKET_INTELLIGENCE_PATH,
|
|
DEFAULT_QUEUE_REFRESH_SECONDS,
|
|
DEFAULT_THROUGHPUT_WINDOW_HOURS,
|
|
)
|
|
from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool
|
|
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
|
|
from queue_cleanup import cleanup_certain_oom_tasks
|
|
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
|
from submission_claims import DEFAULT_CLAIMS_PATH
|
|
from template_selector import TemplateSelector
|
|
from version import AGENT_VERSION
|
|
|
|
|
|
DEFAULT_POLL_RUNS_DIR = Path("poll_runs")
|
|
DEFAULT_ARCHITECTURE_BLACKLIST_PATH = Path(
|
|
".modelhub_state/architecture_compatibility_blacklist.json"
|
|
)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Continuously poll ModelHub queue slots and refill submissions.")
|
|
parser.add_argument("--daily-target", type=int, default=0, help="Total submissions to aim for per UTC day; 0 means unlimited")
|
|
parser.add_argument("--gpu", help="Single GPU alias or platform name, for example: k100")
|
|
parser.add_argument("--gpus", help="Comma-separated GPU aliases/platform names. Omit to auto-use all safe GPUs.")
|
|
parser.add_argument("--min-downloads", type=int, default=50, help="Minimum ModelScope download threshold")
|
|
parser.add_argument(
|
|
"--history-stats-threshold",
|
|
type=int,
|
|
default=500,
|
|
help="Minimum local ledger records before using history stats to rank submissions",
|
|
)
|
|
parser.add_argument("--read-concurrency", type=int, default=4, help="Concurrency for read-only remote calls")
|
|
parser.add_argument(
|
|
"--submit-concurrency",
|
|
type=int,
|
|
default=0,
|
|
help="Concurrency for task submission calls (0 = auto based on token/client count)",
|
|
)
|
|
parser.add_argument(
|
|
"--max-submits-per-run",
|
|
type=int,
|
|
default=0,
|
|
help="Maximum tasks to submit in one cycle (0 means unlimited)",
|
|
)
|
|
parser.add_argument(
|
|
"--recent-model-reserve-slots",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_RECENT_MODEL_RESERVE_SLOTS", "10")),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--recent-model-days",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_RECENT_MODEL_DAYS", "7")),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--dynamic-old-model-cleanup-reserve-slots",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_DYNAMIC_OLD_MODEL_CLEANUP_RESERVE_SLOTS", "5")),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument("--disable-candidate-preflight", action="store_true", help=argparse.SUPPRESS)
|
|
parser.add_argument("--llm-classifier-endpoint", default=os.getenv("MODELHUB_LLM_CLASSIFIER_ENDPOINT"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--llm-classifier-model", default=os.getenv("MODELHUB_LLM_CLASSIFIER_MODEL"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--llm-classifier-api-key", default=os.getenv("MODELHUB_LLM_CLASSIFIER_API_KEY"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--llm-classifier-timeout-seconds", type=int, default=int(os.getenv("MODELHUB_LLM_CLASSIFIER_TIMEOUT_SECONDS", "20")), help=argparse.SUPPRESS)
|
|
parser.add_argument("--llm-classifier-min-deny-confidence", type=float, default=float(os.getenv("MODELHUB_LLM_CLASSIFIER_MIN_DENY_CONFIDENCE", "0.85")), help=argparse.SUPPRESS)
|
|
parser.add_argument("--llm-classifier-cache-path", default=os.getenv("MODELHUB_LLM_CLASSIFIER_CACHE_PATH", ".modelhub_state/llm_classifications.json"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--max-scan-models", type=int, default=0, help="Hard cap on total scanned models (0 means auto)")
|
|
parser.add_argument(
|
|
"--scan-multiplier",
|
|
type=int,
|
|
default=4,
|
|
help="Multiplier used when auto-deriving scan limit from quota/queue capacity",
|
|
)
|
|
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("--disable-gpu-strategy", action="store_true", help="Disable adaptive 70/30 proven-GPU scheduling")
|
|
parser.add_argument(
|
|
"--disable-market-intelligence",
|
|
action="store_true",
|
|
help="Disable live queue/throughput and public framework statistics",
|
|
)
|
|
parser.add_argument(
|
|
"--gpu-strategy-refresh-submissions",
|
|
type=int,
|
|
default=200,
|
|
help="Recalculate adaptive GPU choices after this many accepted submissions",
|
|
)
|
|
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(
|
|
"--claims-path",
|
|
default=os.getenv("MODELHUB_AGENT_CLAIMS_PATH", str(DEFAULT_CLAIMS_PATH)),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--submission-exclusions-path",
|
|
default=os.getenv("MODELHUB_SUBMISSION_EXCLUSIONS_PATH", ".modelhub_state/submission_exclusions.jsonl"),
|
|
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(
|
|
"--gpu-strategy-state-path",
|
|
default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument("--gpu-strategy-recent-window", type=int, default=1000, help=argparse.SUPPRESS)
|
|
parser.add_argument("--gpu-strategy-min-long-samples", type=int, default=100, help=argparse.SUPPRESS)
|
|
parser.add_argument(
|
|
"--market-intelligence-state-path",
|
|
default=os.getenv("MODELHUB_MARKET_INTELLIGENCE_PATH", str(DEFAULT_MARKET_INTELLIGENCE_PATH)),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--market-queue-refresh-seconds",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_MARKET_QUEUE_REFRESH_SECONDS", str(DEFAULT_QUEUE_REFRESH_SECONDS))),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--market-framework-refresh-seconds",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_MARKET_FRAMEWORK_REFRESH_SECONDS", str(DEFAULT_FRAMEWORK_REFRESH_SECONDS))),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--market-throughput-window-hours",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_MARKET_THROUGHPUT_WINDOW_HOURS", str(DEFAULT_THROUGHPUT_WINDOW_HOURS))),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--market-fetch-workers",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_MARKET_FETCH_WORKERS", str(DEFAULT_FETCH_WORKERS))),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--market-framework-min-samples",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_MARKET_FRAMEWORK_MIN_SAMPLES", str(DEFAULT_FRAMEWORK_MIN_SAMPLES))),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--capacity-state-path",
|
|
default=os.getenv("MODELHUB_CAPACITY_STATE_PATH", str(DEFAULT_CAPACITY_STATE_PATH)),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--capacity-probe-interval-cycles",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_CAPACITY_PROBE_INTERVAL_CYCLES", "3")),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
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=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS)
|
|
parser.add_argument("--hf-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)
|
|
parser.add_argument("--modelscope-token", default=None, help=argparse.SUPPRESS)
|
|
parser.add_argument("--poll-interval-seconds", type=int, default=15, help="Sleep between polling cycles when no slots are available")
|
|
parser.add_argument("--idle-interval-seconds", type=int, default=60, help="Sleep between cycles when a scan submits nothing")
|
|
parser.add_argument("--post-cycle-cooldown-seconds", type=int, default=2, help="Short sleep after a successful cycle")
|
|
parser.add_argument("--max-cycles", type=int, default=0, help="Optional hard stop after N cycles; 0 means run until quota is reached")
|
|
parser.add_argument("--print-stats", action="store_true", help="Load outcomes, sync, print stats report, and exit")
|
|
parser.add_argument("--disable-queue-cleanup", action="store_true", help=argparse.SUPPRESS)
|
|
parser.add_argument(
|
|
"--queue-cleanup-interval-cycles",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_QUEUE_CLEANUP_INTERVAL_CYCLES", "120")),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--queue-cleanup-read-concurrency",
|
|
type=int,
|
|
default=int(os.getenv("MODELHUB_QUEUE_CLEANUP_READ_CONCURRENCY", "6")),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--queue-cleanup-report-path",
|
|
default=os.getenv("MODELHUB_QUEUE_CLEANUP_REPORT_PATH", ".modelhub_state/queue_cleanup_latest.json"),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
parser.add_argument(
|
|
"--architecture-blacklist-path",
|
|
default=os.getenv(
|
|
"MODELHUB_ARCHITECTURE_BLACKLIST_PATH",
|
|
str(DEFAULT_ARCHITECTURE_BLACKLIST_PATH),
|
|
),
|
|
help=argparse.SUPPRESS,
|
|
)
|
|
return parser
|
|
|
|
|
|
def _make_cycle_args(base_args: argparse.Namespace) -> argparse.Namespace:
|
|
cycle_args = argparse.Namespace(**vars(base_args))
|
|
cycle_args.rounds = 1
|
|
# Poll runner handles outcome sync on its own schedule (every OUTCOME_SYNC_INTERVAL cycles).
|
|
# Skip the per-cycle sync inside run_submission to avoid redundant paginated API calls.
|
|
cycle_args.skip_outcome_sync = True
|
|
return cycle_args
|
|
|
|
|
|
def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClientPool:
|
|
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
|
|
token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token]
|
|
clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values]
|
|
return ModelHubClientPool(
|
|
clients,
|
|
capacity_probe_interval_cycles=max(0, int(getattr(base_args, "capacity_probe_interval_cycles", 3) or 0)),
|
|
capacity_state_path=Path(getattr(base_args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)),
|
|
recent_model_reserve_slots=max(0, int(getattr(base_args, "recent_model_reserve_slots", 10) or 0)),
|
|
recent_model_days=max(1, int(getattr(base_args, "recent_model_days", 7) or 7)),
|
|
)
|
|
|
|
|
|
def resolve_age_cleanup_policy(
|
|
base_args: argparse.Namespace,
|
|
*,
|
|
initial_cleanup_pending: bool,
|
|
) -> tuple[str, int]:
|
|
"""Reserve ten slots initially, then use a five-slot cleanup hysteresis."""
|
|
initial_reserve_slots = max(
|
|
0,
|
|
int(getattr(base_args, "recent_model_reserve_slots", 10) or 0),
|
|
)
|
|
dynamic_reserve_slots = min(
|
|
initial_reserve_slots,
|
|
max(
|
|
0,
|
|
int(getattr(base_args, "dynamic_old_model_cleanup_reserve_slots", 5) or 0),
|
|
),
|
|
)
|
|
if initial_cleanup_pending:
|
|
return "initial", initial_reserve_slots
|
|
return "dynamic", dynamic_reserve_slots
|
|
|
|
|
|
def _persist_architecture_blacklist(
|
|
report: dict[str, Any],
|
|
*,
|
|
path: Path,
|
|
) -> set[str]:
|
|
blocks = report.get("architectureCompatibilityBlocks") or {}
|
|
blocks = blocks if isinstance(blocks, dict) else {}
|
|
write_json(
|
|
path,
|
|
{
|
|
"generatedAt": report.get("generatedAt"),
|
|
"summary": report.get("architectureCompatibilitySummary") or {},
|
|
"blocks": blocks,
|
|
},
|
|
)
|
|
return set(str(key) for key in blocks)
|
|
|
|
|
|
def _load_task_compatibility_contexts(
|
|
outcome_tracker: OutcomeTracker,
|
|
*,
|
|
ledger_path: Path,
|
|
) -> dict[str, dict[str, Any]]:
|
|
contexts = outcome_tracker.get_task_compatibility_contexts()
|
|
for record in read_jsonl(ledger_path):
|
|
task_id_value = record.get("taskId")
|
|
if task_id_value is None:
|
|
continue
|
|
task_id = str(task_id_value)
|
|
existing = contexts.get(task_id)
|
|
if existing is None:
|
|
contexts[task_id] = {
|
|
"taskId": task_id,
|
|
"modelId": str(record.get("modelId") or ""),
|
|
"targetGpu": str(record.get("targetGpu") or ""),
|
|
"framework": str(record.get("framework") or ""),
|
|
"taskType": str(record.get("taskType") or ""),
|
|
"modelProfile": {},
|
|
"submitTime": record.get("submitTime"),
|
|
}
|
|
continue
|
|
for field in ("modelId", "targetGpu", "framework", "taskType", "submitTime"):
|
|
if not existing.get(field) and record.get(field):
|
|
existing[field] = record.get(field)
|
|
return contexts
|
|
|
|
|
|
def run_poll_loop(
|
|
*,
|
|
base_args: argparse.Namespace,
|
|
now=None,
|
|
run_fn: Callable[..., dict[str, Any]] = run_daily_batches,
|
|
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()
|
|
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()
|
|
|
|
poll_runs_dir = Path(base_args.poll_runs_dir)
|
|
poll_runs_dir.mkdir(parents=True, exist_ok=True)
|
|
poll_run_dir = make_run_dir(poll_runs_dir, now)
|
|
|
|
outcome_tracker = outcome_tracker or OutcomeTracker(Path(base_args.outcomes_path))
|
|
OUTCOME_SYNC_INTERVAL = 3
|
|
STATS_PRINT_INTERVAL = 10
|
|
|
|
log(f"[poll] version={AGENT_VERSION} poll_run_dir={poll_run_dir}")
|
|
log(
|
|
f"[poll] target={base_args.daily_target} dry_run={str(bool(base_args.dry_run)).lower()} "
|
|
f"poll_interval={base_args.poll_interval_seconds}s idle_interval={base_args.idle_interval_seconds}s"
|
|
)
|
|
|
|
cycle_summaries: list[dict[str, Any]] = []
|
|
queue_cleanup_runs: list[dict[str, Any]] = []
|
|
submitted_total = 0
|
|
cycles = 0
|
|
stopped_reason = "max_cycles_reached"
|
|
initial_age_cleanup_pending = True
|
|
pending_architecture_cleanup = False
|
|
last_cleaned_architecture_blocks: set[str] = set()
|
|
|
|
while True:
|
|
if base_args.max_cycles and cycles >= base_args.max_cycles:
|
|
stopped_reason = "max_cycles_reached"
|
|
break
|
|
|
|
cycles += 1
|
|
if hasattr(modelhub_client, "configure_capacity_probe"):
|
|
modelhub_client.configure_capacity_probe(cycles)
|
|
|
|
outcome_synced_this_cycle = False
|
|
if (
|
|
cycles % OUTCOME_SYNC_INTERVAL == 0
|
|
and not getattr(base_args, "skip_outcome_sync", False)
|
|
):
|
|
try:
|
|
synced = outcome_tracker.sync_from_api(modelhub_client)
|
|
outcome_synced_this_cycle = True
|
|
if synced > 0:
|
|
log(f"[poll] cycle={cycles} outcome_sync_updated={synced}")
|
|
sync_feedback = outcome_tracker.get_stats_report()
|
|
active_block_keys = _persist_architecture_blacklist(
|
|
sync_feedback,
|
|
path=Path(
|
|
getattr(
|
|
base_args,
|
|
"architecture_blacklist_path",
|
|
DEFAULT_ARCHITECTURE_BLACKLIST_PATH,
|
|
)
|
|
),
|
|
)
|
|
new_block_keys = active_block_keys - last_cleaned_architecture_blocks
|
|
if (
|
|
new_block_keys
|
|
and not bool(getattr(base_args, "disable_queue_cleanup", False))
|
|
and isinstance(modelhub_client, ModelHubClientPool)
|
|
):
|
|
pending_architecture_cleanup = True
|
|
log(
|
|
f"[queue-cleanup] dynamic_architecture_blocks_added={len(new_block_keys)} "
|
|
f"cleanup_next=immediate"
|
|
)
|
|
except Exception as exc:
|
|
log(f"[poll] cycle={cycles} outcome_sync_error={exc}")
|
|
|
|
cleanup_interval = max(0, int(getattr(base_args, "queue_cleanup_interval_cycles", 120) or 0))
|
|
scheduled_queue_cleanup = bool(
|
|
cycles == 1 or (cleanup_interval > 0 and cycles % cleanup_interval == 0)
|
|
)
|
|
architecture_only_cleanup = bool(
|
|
pending_architecture_cleanup and not scheduled_queue_cleanup
|
|
)
|
|
should_cleanup_queue = (
|
|
not bool(getattr(base_args, "disable_queue_cleanup", False))
|
|
and isinstance(modelhub_client, ModelHubClientPool)
|
|
and (
|
|
scheduled_queue_cleanup
|
|
or pending_architecture_cleanup
|
|
)
|
|
)
|
|
if should_cleanup_queue:
|
|
try:
|
|
if (
|
|
not outcome_synced_this_cycle
|
|
and not getattr(base_args, "skip_outcome_sync", False)
|
|
):
|
|
synced_before_cleanup = outcome_tracker.sync_from_api(modelhub_client)
|
|
if synced_before_cleanup:
|
|
log(
|
|
f"[queue-cleanup] outcome_sync_updated={synced_before_cleanup}"
|
|
)
|
|
cleanup_feedback = outcome_tracker.get_stats_report()
|
|
cleanup_gpu_memory = cleanup_feedback.get("observedGpuMemoryGiB") or {}
|
|
cleanup_architecture_blocks = (
|
|
cleanup_feedback.get("architectureCompatibilityBlocks") or {}
|
|
)
|
|
active_architecture_block_keys = _persist_architecture_blacklist(
|
|
cleanup_feedback,
|
|
path=Path(
|
|
getattr(
|
|
base_args,
|
|
"architecture_blacklist_path",
|
|
DEFAULT_ARCHITECTURE_BLACKLIST_PATH,
|
|
)
|
|
),
|
|
)
|
|
if active_architecture_block_keys - last_cleaned_architecture_blocks:
|
|
pending_architecture_cleanup = True
|
|
age_cleanup_mode, age_cleanup_reserve_slots = resolve_age_cleanup_policy(
|
|
base_args,
|
|
initial_cleanup_pending=initial_age_cleanup_pending,
|
|
)
|
|
log(
|
|
f"[queue-cleanup] mode={'architecture_dynamic' if architecture_only_cleanup else age_cleanup_mode} "
|
|
f"reserve_recent_slots={age_cleanup_reserve_slots} "
|
|
f"recent_days={max(1, int(getattr(base_args, 'recent_model_days', 7) or 7))}"
|
|
)
|
|
cleanup_summary = cleanup_certain_oom_tasks(
|
|
modelhub_client,
|
|
hf_discovery,
|
|
dry_run=bool(base_args.dry_run),
|
|
read_concurrency=max(1, int(getattr(base_args, "queue_cleanup_read_concurrency", 6) or 6)),
|
|
gpu_memory_gib=cleanup_gpu_memory if isinstance(cleanup_gpu_memory, dict) else None,
|
|
architecture_compatibility_blocks=(
|
|
cleanup_architecture_blocks
|
|
if isinstance(cleanup_architecture_blocks, dict)
|
|
else None
|
|
),
|
|
task_compatibility_contexts=(
|
|
_load_task_compatibility_contexts(
|
|
outcome_tracker,
|
|
ledger_path=Path(base_args.ledger_path),
|
|
)
|
|
),
|
|
architecture_only=architecture_only_cleanup,
|
|
age_reserved_slots=age_cleanup_reserve_slots,
|
|
log=log,
|
|
)
|
|
policy_cancelled_recorded = outcome_tracker.mark_policy_cancellations(
|
|
cleanup_summary["cancelledTasks"]
|
|
)
|
|
if policy_cancelled_recorded:
|
|
outcome_tracker.save()
|
|
cleanup_summary["policyCancelledRecorded"] = policy_cancelled_recorded
|
|
write_json(
|
|
Path(
|
|
getattr(
|
|
base_args,
|
|
"queue_cleanup_report_path",
|
|
".modelhub_state/queue_cleanup_latest.json",
|
|
)
|
|
),
|
|
cleanup_summary,
|
|
)
|
|
queue_cleanup_runs.append(
|
|
{
|
|
"cycle": cycles,
|
|
"mode": (
|
|
"architecture_dynamic"
|
|
if architecture_only_cleanup
|
|
else age_cleanup_mode
|
|
),
|
|
"ageReservedSlots": age_cleanup_reserve_slots,
|
|
"ageQueueThresholds": cleanup_summary["oldModelQueueThresholds"],
|
|
"activeScanned": cleanup_summary["activeScanned"],
|
|
"certainOomCount": cleanup_summary["certainOomCount"],
|
|
"architectureBlockCount": cleanup_summary[
|
|
"architectureBlockCount"
|
|
],
|
|
"architectureIncompatibleCount": cleanup_summary[
|
|
"architectureIncompatibleCount"
|
|
],
|
|
"oldOverflowCount": cleanup_summary["oldOverflowCount"],
|
|
"cancelledCount": cleanup_summary["cancelledCount"],
|
|
"policyCancelledRecorded": policy_cancelled_recorded,
|
|
"stopErrorCount": len(cleanup_summary["stopErrors"]),
|
|
}
|
|
)
|
|
if not architecture_only_cleanup:
|
|
initial_age_cleanup_pending = False
|
|
pending_architecture_cleanup = False
|
|
last_cleaned_architecture_blocks = active_architecture_block_keys
|
|
except Exception as exc:
|
|
log(f"[queue-cleanup] error={type(exc).__name__}: {exc} continue_polling=true")
|
|
|
|
active_counts = modelhub_client.active_task_counts() if hasattr(modelhub_client, "active_task_counts") else []
|
|
capacity_limits = modelhub_client.account_capacity_limits() if hasattr(modelhub_client, "account_capacity_limits") else []
|
|
capacity_probe = modelhub_client.capacity_probe_enabled() if hasattr(modelhub_client, "capacity_probe_enabled") else False
|
|
available_slots = modelhub_client.available_submit_slots() if hasattr(modelhub_client, "available_submit_slots") else None
|
|
log(
|
|
f"[poll] cycle={cycles} active_counts={','.join(str(count) for count in active_counts) if active_counts else 'n/a'} "
|
|
f"capacity_limits={','.join(str(limit) for limit in capacity_limits) if capacity_limits else 'n/a'} "
|
|
f"capacity_probe={'on' if capacity_probe else 'off'} "
|
|
f"available_slots={available_slots if available_slots is not None else 'n/a'}"
|
|
)
|
|
|
|
if available_slots is not None and available_slots <= 0:
|
|
log(f"[poll] cycle={cycles} sleep={base_args.poll_interval_seconds}s reason=no_available_slots")
|
|
time.sleep(base_args.poll_interval_seconds)
|
|
continue
|
|
|
|
cycle_args = _make_cycle_args(base_args)
|
|
cycle_args.capacity_probe_already_configured = True
|
|
cycle_summary = run_fn(
|
|
base_args=cycle_args,
|
|
now=utc_now(),
|
|
hf_discovery=hf_discovery,
|
|
modelhub_client=modelhub_client,
|
|
template_selector=template_selector,
|
|
outcome_tracker=outcome_tracker,
|
|
)
|
|
cycle_summaries.append(cycle_summary)
|
|
submitted_total += cycle_summary["submittedTotal"]
|
|
remaining_before_run = cycle_summary.get("remainingDailyQuotaBeforeRun")
|
|
|
|
log(
|
|
f"[poll] cycle_done submitted_total={cycle_summary['submittedTotal']} "
|
|
f"duplicates={cycle_summary.get('duplicateTotal', 0)} "
|
|
f"remaining_before_run={remaining_before_run if remaining_before_run is not None else 'n/a'} "
|
|
f"stop={cycle_summary['stoppedReason']}"
|
|
)
|
|
|
|
if base_args.daily_target > 0 and remaining_before_run is not None and remaining_before_run <= 0:
|
|
stopped_reason = "daily_target_already_reached"
|
|
break
|
|
|
|
if cycle_summary["submittedTotal"] <= 0:
|
|
duplicate_total = int(cycle_summary.get("duplicateTotal", 0) or 0)
|
|
if duplicate_total > 0 and (available_slots is None or available_slots > 0):
|
|
retry_delay = max(1, int(getattr(base_args, "post_cycle_cooldown_seconds", 2) or 2))
|
|
log(f"[poll] cycle={cycles} sleep={retry_delay}s reason=duplicates_need_replacement_candidates")
|
|
time.sleep(retry_delay)
|
|
else:
|
|
log(f"[poll] cycle={cycles} sleep={base_args.idle_interval_seconds}s reason=no_new_submissions")
|
|
time.sleep(base_args.idle_interval_seconds)
|
|
continue
|
|
|
|
if cycles % STATS_PRINT_INTERVAL == 0:
|
|
try:
|
|
stats = outcome_tracker.get_stats_report()
|
|
totals = stats.get("totals", {})
|
|
log(
|
|
f"[poll] cycle={cycles} outcome_stats "
|
|
f"terminal={totals.get('total', 0)} "
|
|
f"success={totals.get('successCount', 0)} "
|
|
f"failed={totals.get('failureCount', 0)} "
|
|
f"success_rate={totals.get('successRate', 0):.3f} "
|
|
f"failure_rate={totals.get('failureRate', 0):.3f} "
|
|
f"architecture_blocks={len(stats.get('architectureCompatibilityBlocks') or {})}"
|
|
)
|
|
except Exception as exc:
|
|
log(f"[poll] cycle={cycles} outcome_stats_error={exc}")
|
|
|
|
if hasattr(modelhub_client, "available_submit_slots") and modelhub_client.available_submit_slots() <= 0:
|
|
log(f"[poll] cycle={cycles} sleep={base_args.poll_interval_seconds}s reason=queue_refilled")
|
|
time.sleep(base_args.poll_interval_seconds)
|
|
else:
|
|
cooldown = getattr(base_args, "post_cycle_cooldown_seconds", 1)
|
|
if cooldown > 0:
|
|
time.sleep(cooldown)
|
|
|
|
try:
|
|
outcome_tracker.sync_from_api(modelhub_client)
|
|
except Exception:
|
|
pass
|
|
outcome_tracker.save()
|
|
try:
|
|
stats_report = outcome_tracker.get_stats_report()
|
|
except Exception:
|
|
stats_report = {}
|
|
|
|
summary = {
|
|
"generatedAt": now.isoformat(),
|
|
"dryRun": bool(base_args.dry_run),
|
|
"dailyTarget": base_args.daily_target,
|
|
"unlimitedDailyTarget": base_args.daily_target <= 0,
|
|
"cycles": cycles,
|
|
"submittedTotal": submitted_total,
|
|
"stoppedReason": stopped_reason,
|
|
"pollRunDir": str(poll_run_dir),
|
|
"cycleSummaries": cycle_summaries,
|
|
"queueCleanupRuns": queue_cleanup_runs,
|
|
"outcomeStats": stats_report,
|
|
}
|
|
write_json(poll_run_dir / "summary.json", summary)
|
|
log(f"[poll] finished submitted_total={submitted_total} cycles={cycles} stopped_reason={stopped_reason}")
|
|
return summary
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
ensure_tokens(args)
|
|
|
|
if getattr(args, "print_stats", False):
|
|
try:
|
|
modelhub_client = _build_modelhub_client(args)
|
|
tracker = OutcomeTracker(Path(args.outcomes_path))
|
|
tracker.sync_from_api(modelhub_client)
|
|
report = tracker.get_stats_report()
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
except Exception as exc:
|
|
print(f"Failed to generate stats report: {exc}", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
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)
|
|
print(f"poll_run_dir={summary['pollRunDir']}")
|
|
print(f"submitted_total={summary['submittedTotal']}")
|
|
print(f"cycles={summary['cycles']}")
|
|
print(f"stopped_reason={summary['stoppedReason']}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|