832 lines
33 KiB
Python
832 lines
33 KiB
Python
from __future__ import annotations
|
||
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import unittest
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from unittest.mock import patch
|
||
|
||
|
||
PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
|
||
if str(PACKAGE_DIR) in sys.path:
|
||
sys.path.remove(str(PACKAGE_DIR))
|
||
sys.path.insert(0, str(PACKAGE_DIR))
|
||
|
||
from common import read_json, read_jsonl # noqa: E402
|
||
from main import build_adaptive_scan_stages, build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402
|
||
from modelhub_client import ( # noqa: E402
|
||
ModelHubAPIError,
|
||
ModelHubClientPool,
|
||
OldModelQueuePolicyError,
|
||
is_duplicate_submission_error,
|
||
)
|
||
from models import HFModelSummary, ModelInspection # noqa: E402
|
||
from outcome_tracker import OutcomeTracker # noqa: E402
|
||
from submission_claims import SubmissionClaimStore, candidate_key # noqa: E402
|
||
from template_selector import TemplateSelector # noqa: E402
|
||
|
||
|
||
class FakeClient:
|
||
def __init__(self, active_count: int = 0, *, reject_capacity: bool = False) -> None:
|
||
self.active_count = active_count
|
||
self.reject_capacity = reject_capacity
|
||
self.submitted: list[dict] = []
|
||
self.count_calls = 0
|
||
self.search_calls = 0
|
||
self._lock = threading.Lock()
|
||
|
||
def count_active_tasks(self, **_kwargs) -> int: # noqa: ANN003
|
||
with self._lock:
|
||
self.count_calls += 1
|
||
return self.active_count
|
||
|
||
def add_task(self, payload: dict) -> dict:
|
||
if self.reject_capacity:
|
||
raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
|
||
with self._lock:
|
||
self.submitted.append(payload)
|
||
return {"code": 0, "data": {"id": len(self.submitted)}}
|
||
|
||
def search_by_model_id(self, _model_id: str) -> dict:
|
||
with self._lock:
|
||
self.search_calls += 1
|
||
return {"code": 0, "data": {"verifyResult": {}}}
|
||
|
||
def list_tasks(self, **_kwargs) -> list[dict]: # noqa: ANN003
|
||
return []
|
||
|
||
def list_tasks_page(self, **_kwargs) -> dict: # noqa: ANN003
|
||
return {"code": 0, "data": {"records": [], "pages": 0}}
|
||
|
||
def find_recent_task_id(self, *_args, **_kwargs): # noqa: ANN002, ANN003
|
||
return None
|
||
|
||
|
||
class FakeDiscovery:
|
||
def __init__(self, count: int) -> None:
|
||
self.models = [
|
||
HFModelSummary(
|
||
repo_id=f"owner/model-{index}",
|
||
downloads=100,
|
||
last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||
pipeline_tag="text-generation",
|
||
)
|
||
for index in range(count)
|
||
]
|
||
|
||
def list_recent_models(self, **_kwargs): # noqa: ANN003
|
||
return self.models
|
||
|
||
@staticmethod
|
||
def inspect_model(model: HFModelSummary) -> ModelInspection:
|
||
return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"])
|
||
|
||
|
||
class DuplicateThenSuccessClient:
|
||
def __init__(self) -> None:
|
||
self.calls = 0
|
||
self.active = 0
|
||
|
||
@staticmethod
|
||
def begin_cycle() -> None:
|
||
return None
|
||
|
||
def available_submit_slots(self) -> int:
|
||
return max(0, 2 - self.active)
|
||
|
||
@staticmethod
|
||
def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003
|
||
return {"code": 0, "data": {"records": [], "pages": 0}}
|
||
|
||
@staticmethod
|
||
def processed_gpus_for_model(_model_id: str) -> set[str]:
|
||
return set()
|
||
|
||
def add_task(self, _payload: dict) -> dict:
|
||
self.calls += 1
|
||
if self.calls <= 2:
|
||
raise ModelHubAPIError("模型正在验证中，请勿重复提交")
|
||
self.active += 1
|
||
return {"code": 0, "data": {"id": f"task-{self.calls}"}}
|
||
|
||
|
||
class DynamicCapacityClient:
|
||
def __init__(self, *, active: int, limit: int, token: str = "dynamic-test") -> None:
|
||
self.active = active
|
||
self.limit = limit
|
||
self.token = token
|
||
|
||
def count_active_tasks(self, *, max_count: int, **_kwargs) -> int: # noqa: ANN003
|
||
return min(self.active, max_count)
|
||
|
||
def add_task(self, _payload: dict) -> dict:
|
||
if self.active >= self.limit:
|
||
raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
|
||
self.active += 1
|
||
return {"code": 0, "data": {"id": self.active}}
|
||
|
||
|
||
class AutoStrategyClient:
|
||
def __init__(self, available: int) -> None:
|
||
self.available = available
|
||
self.calls = 0
|
||
|
||
@staticmethod
|
||
def begin_cycle() -> None:
|
||
return None
|
||
|
||
def available_submit_slots(self) -> int:
|
||
return self.available
|
||
|
||
@staticmethod
|
||
def list_tasks(**_kwargs) -> list[dict]: # noqa: ANN003
|
||
return []
|
||
|
||
@staticmethod
|
||
def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003
|
||
return {"code": 0, "data": {"records": [], "pages": 0}}
|
||
|
||
@staticmethod
|
||
def processed_gpus_for_model(_model_id: str) -> set[str]:
|
||
return set()
|
||
|
||
def add_task(self, _payload: dict) -> dict:
|
||
self.calls += 1
|
||
if self.calls == 1:
|
||
raise ModelHubAPIError("模型正在验证中，请勿重复提交")
|
||
self.available -= 1
|
||
return {"code": 0, "data": {"id": f"strategy-{self.calls}"}}
|
||
|
||
|
||
class ExhaustedRecentDiscovery:
|
||
def __init__(self) -> None:
|
||
self.calls: list[object] = []
|
||
self.recent = HFModelSummary(
|
||
repo_id="owner/recent-exhausted",
|
||
downloads=100,
|
||
last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||
pipeline_tag="text-generation",
|
||
)
|
||
self.older = [
|
||
HFModelSummary(
|
||
repo_id=f"owner/older-{index}",
|
||
downloads=100,
|
||
last_modified=datetime(2025, 1, 1, tzinfo=timezone.utc),
|
||
pipeline_tag="text-generation",
|
||
)
|
||
for index in range(3)
|
||
]
|
||
|
||
def list_recent_models(self, *, updated_after=None, **_kwargs): # noqa: ANN003
|
||
self.calls.append(updated_after)
|
||
return [self.recent] if updated_after is not None else [self.recent, *self.older]
|
||
|
||
@staticmethod
|
||
def inspect_model(model: HFModelSummary) -> ModelInspection:
|
||
return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"])
|
||
|
||
|
||
class ExhaustedRecentClient(AutoStrategyClient):
|
||
def __init__(self, available: int, processed_gpus: set[str]) -> None:
|
||
super().__init__(available)
|
||
self.processed_gpus = processed_gpus
|
||
|
||
def processed_gpus_for_model(self, model_id: str) -> set[str]:
|
||
if model_id == "owner/recent-exhausted":
|
||
return set(self.processed_gpus)
|
||
return set()
|
||
|
||
def add_task(self, _payload: dict) -> dict:
|
||
self.calls += 1
|
||
self.available -= 1
|
||
return {"code": 0, "data": {"id": f"fallback-{self.calls}"}}
|
||
|
||
|
||
def make_candidate(index: int) -> dict:
|
||
return {
|
||
"repoId": f"owner/model-{index}",
|
||
"modelAddress": f"https://modelscope.cn/models/owner/model-{index}",
|
||
"targetGpu": "NVIDIA-A100",
|
||
}
|
||
|
||
|
||
class ClientPoolConcurrencyTests(unittest.TestCase):
|
||
def test_old_models_use_another_account_and_reserve_last_five_positions(self) -> None:
|
||
below_threshold = FakeClient(active_count=94)
|
||
at_threshold = FakeClient(active_count=95)
|
||
pool = ModelHubClientPool(
|
||
[below_threshold, at_threshold], # type: ignore[list-item]
|
||
active_task_cap=100,
|
||
active_counts_ttl=60,
|
||
recent_model_reserve_slots=5,
|
||
recent_model_days=7,
|
||
instance_id="old-model-threshold-test",
|
||
)
|
||
submitted_at = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||
|
||
pool.add_task_for_model(
|
||
{"model": "old-allowed-as-position-95"},
|
||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||
submitted_at=submitted_at,
|
||
)
|
||
|
||
self.assertEqual(1, len(below_threshold.submitted))
|
||
self.assertEqual(0, len(at_threshold.submitted))
|
||
self.assertEqual(0, pool.old_model_submit_slots())
|
||
with self.assertRaises(OldModelQueuePolicyError):
|
||
pool.add_task_for_model(
|
||
{"model": "old-rejected"},
|
||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||
submitted_at=submitted_at,
|
||
)
|
||
|
||
pool.add_task_for_model(
|
||
{"model": "recent-allowed"},
|
||
model_last_modified=datetime(2026, 8, 10, tzinfo=timezone.utc),
|
||
submitted_at=submitted_at,
|
||
)
|
||
self.assertEqual(2, len(below_threshold.submitted) + len(at_threshold.submitted))
|
||
|
||
def test_concurrent_old_model_submissions_cannot_enter_reserved_five_slots(self) -> None:
|
||
client = FakeClient(active_count=93)
|
||
pool = ModelHubClientPool(
|
||
[client], # type: ignore[list-item]
|
||
active_task_cap=100,
|
||
active_counts_ttl=60,
|
||
recent_model_reserve_slots=5,
|
||
recent_model_days=7,
|
||
)
|
||
submitted_at = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||
|
||
def submit(index: int) -> str:
|
||
try:
|
||
pool.add_task_for_model(
|
||
{"model": f"old-{index}"},
|
||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||
submitted_at=submitted_at,
|
||
)
|
||
return "submitted"
|
||
except OldModelQueuePolicyError:
|
||
return "policy_skipped"
|
||
|
||
with ThreadPoolExecutor(max_workers=8) as executor:
|
||
outcomes = list(executor.map(submit, range(8)))
|
||
|
||
self.assertEqual(2, outcomes.count("submitted"))
|
||
self.assertEqual(6, outcomes.count("policy_skipped"))
|
||
self.assertEqual(2, len(client.submitted))
|
||
self.assertEqual(0, pool.old_model_submit_slots())
|
||
|
||
def test_scan_does_not_expand_beyond_recent_window_when_old_slots_are_full(self) -> None:
|
||
now = datetime(2026, 8, 11, tzinfo=timezone.utc)
|
||
stages = build_adaptive_scan_stages(
|
||
now=now,
|
||
initial_updated_after=now - timedelta(hours=48),
|
||
initial_limit=100,
|
||
allow_older_than_recent_window=False,
|
||
recent_model_days=7,
|
||
)
|
||
|
||
self.assertEqual(["configured_window", "last_7_days"], [stage["name"] for stage in stages])
|
||
self.assertTrue(all(stage["updatedAfter"] >= now - timedelta(days=7) for stage in stages))
|
||
|
||
def test_submit_candidate_reports_old_model_policy_skip_at_dynamic_threshold(self) -> None:
|
||
client = FakeClient(active_count=95)
|
||
pool = ModelHubClientPool(
|
||
[client], # type: ignore[list-item]
|
||
active_task_cap=100,
|
||
recent_model_reserve_slots=5,
|
||
recent_model_days=7,
|
||
)
|
||
result = submit_candidate(
|
||
{
|
||
"repoId": "owner/old",
|
||
"modelAddress": "https://modelscope.cn/models/owner/old",
|
||
"targetGpu": "Iluvatar_bi-150",
|
||
"taskType": "text-generation",
|
||
"framework": "vllm",
|
||
"configParams": "framework: vllm",
|
||
"lastModified": "2025-01-01T00:00:00+00:00",
|
||
},
|
||
pool, # type: ignore[arg-type]
|
||
)
|
||
|
||
self.assertEqual("age_policy_deferred", result["outcome"])
|
||
self.assertEqual("age_policy_skipped", result["reason"])
|
||
self.assertEqual([], client.submitted)
|
||
|
||
def test_submit_capacity_exhaustion_is_deferred_not_failed(self) -> None:
|
||
class CapacityClient:
|
||
@staticmethod
|
||
def add_task(_payload): # noqa: ANN001, ANN205
|
||
raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限（100）")
|
||
|
||
result = submit_candidate(
|
||
{
|
||
"repoId": "owner/model",
|
||
"modelAddress": "https://modelscope.cn/models/owner/model",
|
||
"targetGpu": "gpu-a",
|
||
"framework": "vllm",
|
||
"taskType": "text-generation",
|
||
"configParams": "safe",
|
||
},
|
||
CapacityClient(), # type: ignore[arg-type]
|
||
)
|
||
self.assertEqual("capacity_deferred", result["outcome"])
|
||
self.assertEqual("account_capacity_saturated", result["reason"])
|
||
|
||
def test_transformers_platform_prerequisite_is_deferred_not_failed(self) -> None:
|
||
class PrerequisiteClient:
|
||
@staticmethod
|
||
def add_task(_payload): # noqa: ANN001, ANN205
|
||
raise ModelHubAPIError("该模型必须在非transformers框架验证失败后才可以开启transformers框架验证任务")
|
||
|
||
result = submit_candidate(
|
||
{
|
||
"repoId": "owner/model",
|
||
"modelAddress": "https://modelscope.cn/models/owner/model",
|
||
"targetGpu": "gpu-a",
|
||
"framework": "transformers",
|
||
"taskType": "text-generation",
|
||
"configParams": "safe",
|
||
},
|
||
PrerequisiteClient(), # type: ignore[arg-type]
|
||
)
|
||
self.assertEqual("framework_prerequisite_deferred", result["outcome"])
|
||
|
||
def test_age_policy_deferred_candidate_is_skipped_not_failed(self) -> None:
|
||
class AdmissionRaceClient:
|
||
@staticmethod
|
||
def begin_cycle() -> None:
|
||
return None
|
||
|
||
@staticmethod
|
||
def old_model_submit_slots() -> int:
|
||
return 1
|
||
|
||
@staticmethod
|
||
def old_model_queue_thresholds() -> list[int]:
|
||
return [95]
|
||
|
||
@staticmethod
|
||
def available_submit_slots() -> int:
|
||
return 1
|
||
|
||
@staticmethod
|
||
def processed_gpus_for_model(_model_id: str) -> set[str]:
|
||
return set()
|
||
|
||
@staticmethod
|
||
def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003
|
||
return {"code": 0, "data": {"records": [], "pages": 0}}
|
||
|
||
@staticmethod
|
||
def add_task_for_model(*_args, **_kwargs) -> dict: # noqa: ANN002, ANN003
|
||
raise OldModelQueuePolicyError("old-model capacity filled during submission")
|
||
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
root = Path(temporary_dir)
|
||
args = build_parser().parse_args(
|
||
[
|
||
"--gpus",
|
||
"Iluvatar_bi-150",
|
||
"--task-types",
|
||
"text-generation",
|
||
"--limit",
|
||
"1",
|
||
"--max-scan-models",
|
||
"1",
|
||
"--skip-outcome-sync",
|
||
"--skip-history-archive",
|
||
"--disable-candidate-preflight",
|
||
"--disable-gpu-strategy",
|
||
"--disable-market-intelligence",
|
||
]
|
||
)
|
||
args.runs_dir = str(root / "runs")
|
||
args.ledger_path = str(root / "ledger.jsonl")
|
||
args.outcomes_path = str(root / "outcomes.jsonl")
|
||
args.claims_path = str(root / "claims.jsonl")
|
||
args.history_archive_path = str(root / "history.jsonl")
|
||
args.submission_exclusions_path = str(root / "exclusions.jsonl")
|
||
client = AdmissionRaceClient()
|
||
|
||
summary = run_submission(
|
||
args,
|
||
now=datetime(2026, 2, 1, tzinfo=timezone.utc),
|
||
hf_discovery=FakeDiscovery(1), # type: ignore[arg-type]
|
||
modelhub_client=client, # type: ignore[arg-type]
|
||
template_selector=TemplateSelector(),
|
||
)
|
||
|
||
self.assertEqual(0, summary["submittedCount"])
|
||
self.assertEqual(0, summary["failedCount"])
|
||
self.assertEqual(1, summary["skipReasonCounts"]["age_policy_skipped"])
|
||
|
||
def test_old_model_threshold_tracks_capacity_increase_from_100_to_200(self) -> None:
|
||
client = DynamicCapacityClient(active=100, limit=200)
|
||
pool = ModelHubClientPool(
|
||
[client], # type: ignore[list-item]
|
||
active_task_cap=100,
|
||
recent_model_reserve_slots=5,
|
||
capacity_probe_interval_cycles=1,
|
||
capacity_state_path=None,
|
||
)
|
||
|
||
self.assertEqual([100], pool.active_task_counts())
|
||
self.assertEqual([95], pool.old_model_queue_thresholds())
|
||
pool.observe_capacity_lower_bounds([200])
|
||
|
||
self.assertEqual([200], pool.account_capacity_limits())
|
||
self.assertEqual([195], pool.old_model_queue_thresholds())
|
||
|
||
def test_old_model_admission_fails_closed_when_active_count_is_unknown(self) -> None:
|
||
class UnreadableCountClient(FakeClient):
|
||
def count_active_tasks(self, **_kwargs) -> int: # noqa: ANN003
|
||
raise TimeoutError("count unavailable")
|
||
|
||
client = UnreadableCountClient(active_count=0)
|
||
pool = ModelHubClientPool(
|
||
[client], # type: ignore[list-item]
|
||
active_task_cap=100,
|
||
recent_model_reserve_slots=5,
|
||
capacity_state_path=None,
|
||
)
|
||
|
||
with self.assertRaises(OldModelQueuePolicyError):
|
||
pool.add_task_for_model(
|
||
{"model": "old-deferred"},
|
||
model_last_modified=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
||
submitted_at=datetime(2026, 8, 11, tzinfo=timezone.utc),
|
||
)
|
||
|
||
self.assertEqual([None], pool.old_model_queue_thresholds())
|
||
self.assertEqual([], client.submitted)
|
||
|
||
def test_online_submission_does_not_construct_llm_even_when_key_is_present(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
root = Path(temporary_dir)
|
||
args = build_parser().parse_args(
|
||
[
|
||
"--gpus",
|
||
"Iluvatar_bi-150",
|
||
"--task-types",
|
||
"text-generation",
|
||
"--limit",
|
||
"1",
|
||
"--max-scan-models",
|
||
"1",
|
||
"--skip-outcome-sync",
|
||
"--skip-history-archive",
|
||
]
|
||
)
|
||
args.runs_dir = str(root / "runs")
|
||
args.ledger_path = str(root / "ledger.jsonl")
|
||
args.outcomes_path = str(root / "outcomes.jsonl")
|
||
args.claims_path = str(root / "claims.jsonl")
|
||
args.history_archive_path = str(root / "history.jsonl")
|
||
|
||
with (
|
||
patch.dict(os.environ, {"MODELHUB_QWEN_API_KEY": "must-not-be-used"}),
|
||
patch(
|
||
"llm_classifier.LLMAssistedClassifier.__init__",
|
||
side_effect=AssertionError("online submission constructed an LLM client"),
|
||
),
|
||
):
|
||
summary = run_submission(
|
||
args,
|
||
now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
|
||
hf_discovery=FakeDiscovery(1), # type: ignore[arg-type]
|
||
modelhub_client=AutoStrategyClient(available=1), # type: ignore[arg-type]
|
||
)
|
||
|
||
self.assertFalse(summary["candidatePreflight"]["llm"]["enabled"])
|
||
|
||
def test_exhausted_recent_window_expands_to_older_models_in_same_run(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
root = Path(temporary_dir)
|
||
args = build_parser().parse_args(
|
||
[
|
||
"--task-types",
|
||
"text-generation",
|
||
"--since-hours",
|
||
"48",
|
||
"--limit",
|
||
"4",
|
||
"--max-scan-models",
|
||
"4",
|
||
"--skip-outcome-sync",
|
||
"--skip-history-archive",
|
||
]
|
||
)
|
||
args.runs_dir = str(root / "runs")
|
||
args.ledger_path = str(root / "ledger.jsonl")
|
||
args.outcomes_path = str(root / "outcomes.jsonl")
|
||
args.claims_path = str(root / "claims.jsonl")
|
||
args.history_archive_path = str(root / "history.jsonl")
|
||
args.gpu_strategy_state_path = str(root / "strategy.json")
|
||
selector = TemplateSelector()
|
||
supported_gpus = set(selector.supported_target_gpus("text-generation", auto_only=True))
|
||
discovery = ExhaustedRecentDiscovery()
|
||
client = ExhaustedRecentClient(available=2, processed_gpus=supported_gpus)
|
||
|
||
summary = run_submission(
|
||
args,
|
||
now=datetime(2026, 1, 2, 12, tzinfo=timezone.utc),
|
||
hf_discovery=discovery, # type: ignore[arg-type]
|
||
modelhub_client=client, # type: ignore[arg-type]
|
||
template_selector=selector,
|
||
)
|
||
|
||
self.assertEqual(2, summary["submittedCount"])
|
||
self.assertEqual(len(supported_gpus), summary["skipReasonCounts"]["already_processed_for_gpu"])
|
||
self.assertEqual("all_history", summary["scanStages"][-1]["name"])
|
||
self.assertIsNone(discovery.calls[-1])
|
||
|
||
def test_adaptive_strategy_counts_only_platform_accepted_tasks(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
root = Path(temporary_dir)
|
||
args = build_parser().parse_args(
|
||
[
|
||
"--task-types",
|
||
"text-generation",
|
||
"--limit",
|
||
"4",
|
||
"--max-scan-models",
|
||
"4",
|
||
"--skip-outcome-sync",
|
||
"--skip-history-archive",
|
||
]
|
||
)
|
||
args.runs_dir = str(root / "runs")
|
||
args.ledger_path = str(root / "ledger.jsonl")
|
||
args.outcomes_path = str(root / "outcomes.jsonl")
|
||
args.claims_path = str(root / "claims.jsonl")
|
||
args.history_archive_path = str(root / "history.jsonl")
|
||
args.gpu_strategy_state_path = str(root / "strategy.json")
|
||
client = AutoStrategyClient(available=3)
|
||
|
||
summary = run_submission(
|
||
args,
|
||
now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
|
||
hf_discovery=FakeDiscovery(4), # type: ignore[arg-type]
|
||
modelhub_client=client, # type: ignore[arg-type]
|
||
)
|
||
|
||
state = read_json(root / "strategy.json")
|
||
self.assertEqual(3, summary["submittedCount"])
|
||
self.assertEqual(1, summary["duplicateCount"])
|
||
self.assertEqual(3, state["acceptedSinceRefresh"])
|
||
self.assertEqual(3, sum(state["acceptedByCategory"].values()))
|
||
|
||
def test_duplicate_batch_is_replaced_until_available_slots_are_filled(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
root = Path(temporary_dir)
|
||
args = build_parser().parse_args(
|
||
[
|
||
"--gpus",
|
||
"Iluvatar_bi-150",
|
||
"--task-types",
|
||
"text-generation",
|
||
"--limit",
|
||
"4",
|
||
"--max-scan-models",
|
||
"4",
|
||
"--skip-outcome-sync",
|
||
"--skip-history-archive",
|
||
]
|
||
)
|
||
args.runs_dir = str(root / "runs")
|
||
args.ledger_path = str(root / "ledger.jsonl")
|
||
args.outcomes_path = str(root / "outcomes.jsonl")
|
||
args.claims_path = str(root / "claims.jsonl")
|
||
args.history_archive_path = str(root / "history.jsonl")
|
||
client = DuplicateThenSuccessClient()
|
||
|
||
summary = run_submission(
|
||
args,
|
||
now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
|
||
hf_discovery=FakeDiscovery(4), # type: ignore[arg-type]
|
||
modelhub_client=client, # type: ignore[arg-type]
|
||
)
|
||
|
||
self.assertEqual(2, summary["targetSubmitCount"])
|
||
self.assertEqual(4, summary["plannedSubmitCount"])
|
||
self.assertEqual(2, summary["duplicateCount"])
|
||
self.assertEqual(2, summary["submittedCount"])
|
||
self.assertEqual(0, client.available_submit_slots())
|
||
|
||
def test_platform_duplicate_is_classified_as_skipped_candidate(self) -> None:
|
||
error = ModelHubAPIError("模型正在验证中，请勿重复提交")
|
||
self.assertTrue(is_duplicate_submission_error(error))
|
||
|
||
class DuplicateClient:
|
||
@staticmethod
|
||
def add_task(_payload): # noqa: ANN001
|
||
raise error
|
||
|
||
candidate = {
|
||
**make_candidate(1),
|
||
"taskType": "text-generation",
|
||
"framework": "vllm",
|
||
"configParams": "framework: vllm",
|
||
}
|
||
result = submit_candidate(candidate, DuplicateClient()) # type: ignore[arg-type]
|
||
self.assertEqual("duplicate", result["outcome"])
|
||
|
||
def test_every_account_reaches_capacity_under_load(self) -> None:
|
||
clients = [FakeClient() for _ in range(12)]
|
||
pool = ModelHubClientPool(
|
||
clients, # type: ignore[arg-type]
|
||
active_task_cap=10,
|
||
active_counts_ttl=60,
|
||
instance_id="full-load-test",
|
||
)
|
||
|
||
with ThreadPoolExecutor(max_workers=24) as executor:
|
||
list(executor.map(lambda index: pool.add_task({"index": index}), range(120)))
|
||
|
||
self.assertEqual([10] * 12, [len(client.submitted) for client in clients])
|
||
self.assertEqual([10] * 12, pool.active_task_counts())
|
||
self.assertEqual(0, pool.available_submit_slots())
|
||
|
||
def test_concurrent_submissions_reserve_and_balance_accounts(self) -> None:
|
||
clients = [FakeClient(), FakeClient()]
|
||
pool = ModelHubClientPool(
|
||
clients, # type: ignore[arg-type]
|
||
active_task_cap=4,
|
||
active_counts_ttl=60,
|
||
instance_id="balance-test",
|
||
)
|
||
|
||
with ThreadPoolExecutor(max_workers=8) as executor:
|
||
results = list(executor.map(lambda index: pool.add_task({"index": index}), range(8)))
|
||
|
||
self.assertEqual(8, len(results))
|
||
self.assertEqual([4, 4], [len(client.submitted) for client in clients])
|
||
self.assertEqual([4, 4], pool.active_task_counts())
|
||
self.assertEqual(0, pool.available_submit_slots())
|
||
# A lagging remote count must not erase successful local reservations.
|
||
pool._refresh_active_counts(force=True)
|
||
self.assertEqual([4, 4], pool.active_task_counts())
|
||
|
||
def test_capacity_rejection_falls_through_to_another_account(self) -> None:
|
||
full_elsewhere = FakeClient(active_count=0, reject_capacity=True)
|
||
available = FakeClient(active_count=1)
|
||
pool = ModelHubClientPool(
|
||
[full_elsewhere, available], # type: ignore[arg-type]
|
||
active_task_cap=2,
|
||
instance_id="fallback-test",
|
||
)
|
||
|
||
response = pool.add_task({"model": "x"})
|
||
|
||
self.assertEqual(1, response["data"]["id"])
|
||
self.assertEqual(0, len(full_elsewhere.submitted))
|
||
self.assertEqual(1, len(available.submitted))
|
||
|
||
def test_verification_cache_is_reused_between_cycles(self) -> None:
|
||
client = FakeClient()
|
||
pool = ModelHubClientPool([client], active_task_cap=2) # type: ignore[arg-type]
|
||
pool.search_by_model_id("owner/model")
|
||
pool.search_by_model_id("owner/model")
|
||
self.assertEqual(1, client.search_calls)
|
||
pool.begin_cycle()
|
||
pool.search_by_model_id("owner/model")
|
||
self.assertEqual(1, client.search_calls)
|
||
|
||
def test_capacity_probe_discovers_a_higher_dynamic_account_limit(self) -> None:
|
||
client = DynamicCapacityClient(active=2, limit=3)
|
||
pool = ModelHubClientPool(
|
||
[client], # type: ignore[arg-type]
|
||
active_task_cap=2,
|
||
active_counts_ttl=60,
|
||
capacity_probe_interval_cycles=3,
|
||
instance_id="capacity-growth-test",
|
||
)
|
||
|
||
pool.configure_capacity_probe(2)
|
||
self.assertEqual(0, pool.available_submit_slots())
|
||
pool.configure_capacity_probe(3)
|
||
self.assertEqual(1, pool.available_submit_slots())
|
||
pool.add_task({"model": "probe"})
|
||
|
||
self.assertEqual([3], pool.account_capacity_limits())
|
||
self.assertEqual([3], pool.active_task_counts())
|
||
self.assertEqual(0, pool.available_submit_slots())
|
||
|
||
def test_rejected_capacity_probe_keeps_known_limit_and_enters_cooldown(self) -> None:
|
||
client = DynamicCapacityClient(active=2, limit=2)
|
||
pool = ModelHubClientPool(
|
||
[client], # type: ignore[arg-type]
|
||
active_task_cap=2,
|
||
active_counts_ttl=60,
|
||
capacity_probe_interval_cycles=3,
|
||
capacity_probe_cooldown_cycles=3,
|
||
instance_id="capacity-rejection-test",
|
||
)
|
||
|
||
pool.configure_capacity_probe(3)
|
||
self.assertEqual(1, pool.available_submit_slots())
|
||
with self.assertRaises(ModelHubAPIError):
|
||
pool.add_task({"model": "probe"})
|
||
|
||
self.assertEqual([2], pool.account_capacity_limits())
|
||
self.assertEqual(0, pool.available_submit_slots())
|
||
|
||
|
||
class ProcessCoordinationTests(unittest.TestCase):
|
||
def test_local_failure_block_expires_after_24_hours(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
tracker = OutcomeTracker(Path(temporary_dir) / "outcomes.jsonl")
|
||
now = datetime(2026, 1, 2, 12, tzinfo=timezone.utc)
|
||
tracker._records = [
|
||
{
|
||
"modelId": "owner/old-failure",
|
||
"targetGpu": "gpu-a",
|
||
"outcome": "failed",
|
||
"failureCategory": "model_load",
|
||
"failureScope": "model_gpu_framework",
|
||
"submitTime": (now - timedelta(hours=26)).isoformat(),
|
||
"lastSyncTime": (now - timedelta(hours=25)).isoformat(),
|
||
},
|
||
{
|
||
"modelId": "owner/recent-failure",
|
||
"targetGpu": "gpu-a",
|
||
"outcome": "failed",
|
||
"failureCategory": "model_load",
|
||
"failureScope": "model_gpu_framework",
|
||
"submitTime": (now - timedelta(hours=2)).isoformat(),
|
||
"lastSyncTime": (now - timedelta(hours=1)).isoformat(),
|
||
},
|
||
]
|
||
tracker._rebuild_indexes()
|
||
|
||
self.assertFalse(tracker.is_model_gpu_failed("owner/old-failure", "gpu-a", now=now))
|
||
self.assertTrue(tracker.is_model_gpu_failed("owner/recent-failure", "gpu-a", now=now))
|
||
|
||
def test_duplicate_claims_are_retained_so_next_batch_moves_forward(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
path = Path(temporary_dir) / "claims.jsonl"
|
||
candidates = [make_candidate(index) for index in range(4)]
|
||
store = SubmissionClaimStore(path, owner_id="worker-a")
|
||
|
||
duplicates = store.claim(candidates, limit=2)
|
||
store.mark_submitted(duplicates)
|
||
replacements = store.claim(candidates, limit=2)
|
||
|
||
self.assertEqual(
|
||
{candidate_key(candidate) for candidate in candidates[2:]},
|
||
{candidate_key(candidate) for candidate in replacements},
|
||
)
|
||
|
||
def test_concurrent_claim_stores_select_disjoint_candidates(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
path = Path(temporary_dir) / "claims.jsonl"
|
||
candidates = [make_candidate(index) for index in range(6)]
|
||
first = SubmissionClaimStore(path, owner_id="worker-a")
|
||
second = SubmissionClaimStore(path, owner_id="worker-b")
|
||
|
||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||
first_future = executor.submit(first.claim, candidates, limit=3)
|
||
second_future = executor.submit(second.claim, candidates, limit=3)
|
||
first_claims = first_future.result()
|
||
second_claims = second_future.result()
|
||
|
||
first_keys = {candidate_key(candidate) for candidate in first_claims}
|
||
second_keys = {candidate_key(candidate) for candidate in second_claims}
|
||
self.assertEqual(3, len(first_claims))
|
||
self.assertEqual(3, len(second_claims))
|
||
self.assertTrue(first_keys.isdisjoint(second_keys))
|
||
|
||
def test_outcome_saves_merge_instead_of_overwriting_other_worker(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
path = Path(temporary_dir) / "outcomes.jsonl"
|
||
first = OutcomeTracker(path)
|
||
second = OutcomeTracker(path)
|
||
first.record_submission("model-a", "gpu", "vllm", "text-generation", "task-a", "2026-01-01T00:00:00+00:00")
|
||
second.record_submission("model-b", "gpu", "vllm", "text-generation", "task-b", "2026-01-01T00:00:01+00:00")
|
||
|
||
first.save()
|
||
second.save()
|
||
|
||
self.assertEqual({"task-a", "task-b"}, {row["taskId"] for row in read_jsonl(path)})
|
||
|
||
def test_run_directories_are_allocated_atomically(self) -> None:
|
||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||
base = Path(temporary_dir) / "runs"
|
||
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||
with ThreadPoolExecutor(max_workers=8) as executor:
|
||
paths = list(executor.map(lambda _index: make_run_dir(base, now), range(8)))
|
||
|
||
self.assertEqual(8, len(set(paths)))
|
||
self.assertTrue(all(path.is_dir() for path in paths))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|