fix: refill from expanded candidate history
This commit is contained in:
CoolBoy
34
tests/test_agent_entrypoint.py
Normal file
34
tests/test_agent_entrypoint.py
Normal file
@@ -0,0 +1,34 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
ROOT_DIR = Path(__file__).resolve().parents[1]
|
||||
SPEC = importlib.util.spec_from_file_location("modelhub_agent_entrypoint", ROOT_DIR / "main.py")
|
||||
if SPEC is None or SPEC.loader is None:
|
||||
raise RuntimeError("Unable to load the hosted agent entrypoint")
|
||||
ENTRYPOINT = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(ENTRYPOINT)
|
||||
|
||||
|
||||
class HostedAgentEntrypointTests(unittest.TestCase):
|
||||
def test_hosted_worker_always_uses_unlimited_cycle_submissions(self) -> None:
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"MODELHUB_AGENT_MAX_SUBMITS_PER_RUN": "5",
|
||||
"MODELHUB_AGENT_EXTRA_ARGS": "--max-submits-per-run 3",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
command = ENTRYPOINT._worker_command()
|
||||
|
||||
self.assertEqual(["--max-submits-per-run", "0"], command[-2:])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -5,7 +5,7 @@ import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_s
|
||||
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:
|
||||
@@ -154,6 +155,50 @@ class AutoStrategyClient:
|
||||
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}",
|
||||
@@ -163,6 +208,47 @@ def make_candidate(index: int) -> dict:
|
||||
|
||||
|
||||
class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
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)
|
||||
@@ -305,7 +391,7 @@ class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
self.assertEqual(0, len(full_elsewhere.submitted))
|
||||
self.assertEqual(1, len(available.submitted))
|
||||
|
||||
def test_verification_cache_is_reset_between_cycles(self) -> None:
|
||||
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")
|
||||
@@ -313,7 +399,7 @@ class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
self.assertEqual(1, client.search_calls)
|
||||
pool.begin_cycle()
|
||||
pool.search_by_model_id("owner/model")
|
||||
self.assertEqual(2, client.search_calls)
|
||||
self.assertEqual(1, client.search_calls)
|
||||
|
||||
def test_capacity_probe_discovers_a_higher_dynamic_account_limit(self) -> None:
|
||||
client = DynamicCapacityClient(active=2, limit=3)
|
||||
@@ -356,6 +442,31 @@ class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
|
||||
|
||||
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",
|
||||
"submitTime": (now - timedelta(hours=26)).isoformat(),
|
||||
"lastSyncTime": (now - timedelta(hours=25)).isoformat(),
|
||||
},
|
||||
{
|
||||
"modelId": "owner/recent-failure",
|
||||
"targetGpu": "gpu-a",
|
||||
"outcome": "failed",
|
||||
"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"
|
||||
|
||||
Reference in New Issue
Block a user