feat: reserve queue capacity for recent models

This commit is contained in:
CoolBoy
2026-08-11 01:29:54 +08:00
parent 38ab25fc3c
commit f12d96b138
13 changed files with 927 additions and 46 deletions

View File

@@ -55,6 +55,9 @@ Optional tuning:
- `MODELHUB_QUEUE_CLEANUP_INTERVAL_CYCLES` default `120`; cleanup also runs once at startup
- `MODELHUB_QUEUE_CLEANUP_READ_CONCURRENCY` default `6`
- `MODELHUB_QUEUE_CLEANUP_REPORT_PATH` default `.modelhub_state/queue_cleanup_latest.json`
- `MODELHUB_OLD_MODEL_QUEUE_THRESHOLD` default `80` per account
- `MODELHUB_DYNAMIC_OLD_MODEL_CLEANUP_THRESHOLD` default `95` per account
- `MODELHUB_RECENT_MODEL_DAYS` default `7`
## Adaptive GPU Strategy
@@ -135,6 +138,25 @@ same model or infer failure from historical similarity. The cleanup repeats
every 120 poll cycles by default and writes its full evidence report to
`.modelhub_state/queue_cleanup_latest.json`.
Each account reserves only its first 80 active queue positions for models older
than seven days. Old-model reservations are made under the same account lock as
capacity reservations, so concurrent submissions cannot cross position 80.
Once every account has at least 80 active tasks, discovery is capped at the
seven-day window and only recent models can fill positions 81100. Unknown model
timestamps are treated as old for new submissions.
Queue cleanup applies a two-level version of the policy. On poller cycle 1 it
first stops deterministic OOM tasks, recalculates each account's surviving task
order by numeric task ID, and removes older-than-seven-days tasks after position
80. On later scheduled cleanup cycles, the age threshold relaxes to 95: old
tasks in positions 81-95 are retained and only old tasks at position 96 or later
are removed. New old-model submissions are still blocked at position 80, so the
upper queue remains available to recent models without repeatedly over-pruning
existing work. Recent overflow tasks are always kept. ModelScope metadata
failures fail closed and never trigger cancellation. Immediately before
mutation, task ownership, active status, and post-OOM queue position are checked
again.
The verified capacities, safe repository-size boundaries, evidence hierarchy,
and source links are recorded in
`docs/gpu-memory-capacity-2026-08-10.md`.
@@ -218,12 +240,16 @@ second active-state check and fail-closed handling for incomplete evidence.
Version `2026.08.11.2` removes LLM inference from every online path; credentials
alone cannot activate it, and unresolved cases remain available for offline,
human-reviewed rule development.
Version `2026.08.11.3` adds atomic per-account 80/7-day admission, prevents
history expansion once old-model positions are exhausted, performs a strict
position-80 startup cleanup, and relaxes scheduled age cleanup to position 95
after OOM cleanup and a second queue check.
## Deploy
Create a tag and submit the repository URL plus tag in "我的适配智能体".
```bash
git tag agent-v12
git push origin agent-v12
git tag agent-v17
git push origin agent-v17
```

View File

@@ -120,6 +120,12 @@ bash run_poll.sh --dry-run
memory rule as new submissions. Only tasks whose own repository size times `1.20` exceeds
their selected GPU capacity are stopped, after a fresh account-scoped active-state check.
Incomplete size/capacity evidence is never used for cancellation.
- Models older than seven days may occupy only positions 180 of each account's active queue.
The account pool enforces the boundary atomically. Once all accounts reach 80, discovery no
longer expands beyond seven days. Cleanup stops OOM tasks first, recalculates the surviving
queue order, and on startup stops old tasks still ranked 81 or later. Scheduled cleanup then
retains positions 8195 and stops only old tasks ranked 96 or later. Recent overflow tasks stay.
Unknown ModelScope timestamps never authorize a cancellation.
## Important Flags
@@ -153,7 +159,12 @@ Common flags:
- `--submit-concurrency`: concurrent task submission calls used by each cycle (0 = auto)
- `--post-cycle-cooldown-seconds`: pause after a successful cycle before next cycle (default 2)
- `--max-cycles`: optional hard stop for testing or batch windows
- `--disable-queue-cleanup`: disable automatic deterministic OOM cleanup
- `--disable-queue-cleanup`: disable automatic OOM and old-overflow queue cleanup
The first automatic cleanup removes models older than seven days after each
account's position 80. Later scheduled cleanups use position 95, while admission
continues to reserve positions 81-100 for recent models. Override the later
threshold with `MODELHUB_DYNAMIC_OLD_MODEL_CLEANUP_THRESHOLD`.
Failure-informed preflight is enabled by default. It rejects deterministic
missing-file and predicted-OOM cases, clamps unsafe context-length arguments,

View File

@@ -85,6 +85,18 @@ def build_parser() -> argparse.ArgumentParser:
default=0,
help="Maximum tasks to submit in one run (0 means unlimited)",
)
parser.add_argument(
"--old-model-queue-threshold",
type=int,
default=int(os.getenv("MODELHUB_OLD_MODEL_QUEUE_THRESHOLD", "80")),
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("--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)
@@ -240,6 +252,8 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar
capacity_probe_interval_cycles=getattr(base_args, "capacity_probe_interval_cycles", 3),
submit_concurrency=getattr(base_args, "submit_concurrency", 1),
max_submits_per_run=getattr(base_args, "max_submits_per_run", 0),
old_model_queue_threshold=getattr(base_args, "old_model_queue_threshold", 80),
recent_model_days=getattr(base_args, "recent_model_days", 7),
disable_candidate_preflight=getattr(base_args, "disable_candidate_preflight", False),
llm_classifier_endpoint=getattr(base_args, "llm_classifier_endpoint", None),
llm_classifier_model=getattr(base_args, "llm_classifier_model", None),

View File

@@ -3,6 +3,7 @@ from __future__ import annotations
import os
import threading
import time
from datetime import datetime, timezone
from pathlib import PurePosixPath
from typing import Any
from urllib.parse import quote
@@ -71,6 +72,8 @@ class HuggingFaceDiscovery:
self._repo_tree_lock = threading.Lock()
self._model_config_cache: dict[str, tuple[dict[str, Any], str | None]] = {}
self._model_config_lock = threading.Lock()
self._model_last_modified_cache: dict[str, datetime | None] = {}
self._model_last_modified_lock = threading.Lock()
self._model_page_cache: dict[tuple[str, int, int], tuple[float, list[dict[str, Any]]]] = {}
self._model_page_cache_ttl = max(
0.0,
@@ -237,6 +240,47 @@ class HuggingFaceDiscovery:
self._model_config_cache[repo_id] = result
return dict(result[0]), result[1]
def get_model_last_modified(self, repo_id: str) -> datetime | None:
with self._model_last_modified_lock:
if repo_id in self._model_last_modified_cache:
return self._model_last_modified_cache[repo_id]
encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
try:
payload = self.legacy_http_client.request_json(
"GET",
f"/api/v1/models/{encoded_repo_id}",
)
except HttpJsonError as exc:
print(f"[modelscope] model_metadata_error repo={repo_id} error={exc}", flush=True)
payload = None
data = payload.get("Data") if isinstance(payload, dict) else None
if not isinstance(data, dict) and isinstance(payload, dict):
data = payload.get("data")
result: datetime | None = None
if isinstance(data, dict):
raw_value = (
data.get("LastUpdatedTime")
or data.get("lastUpdatedTime")
or data.get("last_modified")
or data.get("updated_at")
)
if isinstance(raw_value, (int, float)):
timestamp = float(raw_value)
if timestamp > 10_000_000_000:
timestamp /= 1000.0
try:
result = datetime.fromtimestamp(timestamp, tz=timezone.utc)
except (OverflowError, OSError, ValueError):
result = None
else:
result = parse_datetime(raw_value)
with self._model_last_modified_lock:
self._model_last_modified_cache[repo_id] = result
return result
def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]:
with self._repo_tree_lock:
cached = self._repo_tree_cache.get(repo_id)

View File

@@ -33,6 +33,7 @@ from modelhub_client import (
ModelHubAPIError,
ModelHubClient,
ModelHubClientPool,
OldModelQueuePolicyError,
is_duplicate_submission_error,
is_model_uniqueness_error,
)
@@ -107,6 +108,18 @@ def build_parser() -> argparse.ArgumentParser:
default=0,
help="Maximum tasks to submit in one run (0 means unlimited)",
)
parser.add_argument(
"--old-model-queue-threshold",
type=int,
default=int(os.getenv("MODELHUB_OLD_MODEL_QUEUE_THRESHOLD", "80")),
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(
"--disable-candidate-preflight",
action="store_true",
@@ -548,9 +561,17 @@ def build_adaptive_scan_stages(
initial_updated_after,
initial_limit: int,
max_models: int = ADAPTIVE_SCAN_MAX_MODELS,
allow_older_than_recent_window: bool = True,
recent_model_days: int = 7,
) -> list[dict[str, Any]]:
max_models = max(1, min(ADAPTIVE_SCAN_MAX_MODELS, int(max_models)))
initial_limit = min(max_models, max(1, int(initial_limit)))
recent_cutoff = now - timedelta(days=max(1, int(recent_model_days)))
if not allow_older_than_recent_window and (
initial_updated_after is None or initial_updated_after < recent_cutoff
):
initial_updated_after = recent_cutoff
stages: list[dict[str, Any]] = [
{
"name": "configured_window",
@@ -559,17 +580,16 @@ def build_adaptive_scan_stages(
}
]
seven_days_ago = now - timedelta(days=7)
thirty_days_ago = now - timedelta(days=30)
if initial_updated_after is not None and initial_updated_after > seven_days_ago:
if initial_updated_after is not None and initial_updated_after > recent_cutoff:
stages.append(
{
"name": "last_7_days",
"updatedAfter": seven_days_ago,
"updatedAfter": recent_cutoff,
"limit": min(max_models, max(initial_limit, ADAPTIVE_SCAN_MIN_FALLBACK_MODELS)),
}
)
if initial_updated_after is not None and initial_updated_after > thirty_days_ago:
if allow_older_than_recent_window and initial_updated_after is not None and initial_updated_after > thirty_days_ago:
stages.append(
{
"name": "last_30_days",
@@ -577,7 +597,7 @@ def build_adaptive_scan_stages(
"limit": min(max_models, max(initial_limit, 1500)),
}
)
if initial_updated_after is not None or initial_limit < max_models:
if allow_older_than_recent_window and (initial_updated_after is not None or initial_limit < max_models):
stages.append(
{
"name": "all_history",
@@ -598,6 +618,13 @@ def build_adaptive_scan_stages(
return deduped
def is_recent_model(last_modified: Any, *, reference_time, recent_model_days: int = 7) -> bool:
parsed = parse_datetime(last_modified)
if parsed is None:
return False
return parsed >= reference_time - timedelta(days=max(1, int(recent_model_days)))
def collect_candidates_from_models(
*,
models: list[HFModelSummary],
@@ -702,7 +729,14 @@ def submit_candidate(
payload["strategyId"] = strategy_id
submit_time = utc_now()
try:
response = modelhub_client.add_task(payload)
if hasattr(modelhub_client, "add_task_for_model"):
response = modelhub_client.add_task_for_model(
payload,
model_last_modified=candidate.get("lastModified"),
submitted_at=submit_time,
)
else:
response = modelhub_client.add_task(payload)
task_id = extract_task_id_from_submit_response(response)
if task_id is None:
task_id = modelhub_client.find_recent_task_id(candidate["repoId"], candidate["targetGpu"], submit_time)
@@ -714,6 +748,12 @@ def submit_candidate(
"responseData": response.get("data"),
}
except ModelHubAPIError as exc:
if isinstance(exc, OldModelQueuePolicyError):
return {
"outcome": "old_model_policy_skipped",
"candidate": candidate,
"reason": "all_accounts_at_old_model_queue_threshold",
}
if is_model_uniqueness_error(exc):
print(
f"[submit] skipped repo={candidate['repoId']} gpu={candidate['targetGpu']} "
@@ -799,6 +839,8 @@ def run_submission(
clients,
capacity_probe_interval_cycles=max(0, int(getattr(args, "capacity_probe_interval_cycles", 3) or 0)),
capacity_state_path=Path(getattr(args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)),
old_model_queue_threshold=max(1, int(getattr(args, "old_model_queue_threshold", 80) or 80)),
recent_model_days=max(1, int(getattr(args, "recent_model_days", 7) or 7)),
)
if hasattr(modelhub_client, "begin_cycle"):
@@ -1065,11 +1107,25 @@ def run_submission(
pipeline_tags = pipeline_tags_for_task_types(selected_task_types)
explicit_scan_cap = max(0, int(getattr(args, "max_scan_models", 0) or 0))
recent_model_days = max(1, int(getattr(args, "recent_model_days", 7) or 7))
old_model_queue_threshold = max(1, int(getattr(args, "old_model_queue_threshold", 80) or 80))
old_model_slots_before_scan: int | None = None
if hasattr(modelhub_client, "old_model_submit_slots"):
old_model_slots_before_scan = int(modelhub_client.old_model_submit_slots())
allow_older_models_for_scan = old_model_slots_before_scan is None or old_model_slots_before_scan > 0
print(
f"[age-policy] recent_days={recent_model_days} account_threshold={old_model_queue_threshold} "
f"old_model_slots={old_model_slots_before_scan if old_model_slots_before_scan is not None else 'n/a'} "
f"scan_older={'on' if allow_older_models_for_scan else 'off'}",
flush=True,
)
for stage in build_adaptive_scan_stages(
now=now,
initial_updated_after=updated_after,
initial_limit=scan_limit,
max_models=explicit_scan_cap or ADAPTIVE_SCAN_MAX_MODELS,
allow_older_than_recent_window=allow_older_models_for_scan,
recent_model_days=recent_model_days,
):
if candidate_goal <= 0 or len(candidates) >= candidate_goal:
break
@@ -1088,6 +1144,16 @@ def run_submission(
)
except TypeError:
models = hf_discovery.list_recent_models(**query_kwargs)
if not allow_older_models_for_scan:
models = [
model
for model in models
if is_recent_model(
model.last_modified,
reference_time=now,
recent_model_days=recent_model_days,
)
]
stage_candidates, stage_skipped, stage_failed, processed_count = collect_candidates_from_models(
models=models,
@@ -1165,11 +1231,25 @@ def run_submission(
# already-scanned pool until the desired number of real submissions is
# reached or account capacity is genuinely exhausted.
while len(submitted) < target_submit_count and len(attempted_candidates) < max_submit_attempts:
allow_older_models_now = (
bool(modelhub_client.can_submit_old_models())
if hasattr(modelhub_client, "can_submit_old_models")
else True
)
submission_reference_time = utc_now()
remaining_candidates = [
candidate
for candidate in diversified_candidates
if candidate_key(candidate) not in attempted_keys
and not submission_exclusion_store.is_blocked(candidate["repoId"], candidate["targetGpu"])
and (
allow_older_models_now
or is_recent_model(
candidate.get("lastModified"),
reference_time=submission_reference_time,
recent_model_days=recent_model_days,
)
)
]
remaining_candidates = one_candidate_per_model(remaining_candidates)
desired_count = min(
@@ -1209,6 +1289,7 @@ def run_submission(
batch_submitted_candidates: list[dict[str, Any]] = []
batch_duplicate_candidates: list[dict[str, Any]] = []
batch_uniqueness_rejected_candidates: list[dict[str, Any]] = []
batch_policy_skipped_candidates: list[dict[str, Any]] = []
batch_failed_candidates: list[dict[str, Any]] = []
for index in range(len(batch_candidates)):
result = ordered_results.get(index)
@@ -1243,6 +1324,16 @@ def run_submission(
}
)
continue
if result["outcome"] == "old_model_policy_skipped":
batch_policy_skipped_candidates.append(candidate)
skipped.append(
{
"repoId": candidate["repoId"],
"targetGpu": candidate["targetGpu"],
"reason": "all_accounts_at_old_model_queue_threshold",
}
)
continue
if result["outcome"] in {"failed", "precheck_deferred"}:
batch_failed_candidates.append(candidate)
failed.append(
@@ -1290,7 +1381,7 @@ def run_submission(
claim_store.mark_submitted(
[*batch_submitted_candidates, *batch_duplicate_candidates, *batch_uniqueness_rejected_candidates]
)
claim_store.release(batch_failed_candidates)
claim_store.release([*batch_failed_candidates, *batch_policy_skipped_candidates])
if strategy_manager is not None:
strategy_manager.record_accepted(batch_submitted_candidates)
@@ -1332,6 +1423,12 @@ def run_submission(
"gpuStrategy": strategy_summary,
"marketIntelligence": market_summary,
"candidatePreflight": preflight_summary,
"agePolicy": {
"recentModelDays": recent_model_days,
"oldModelSubmitThreshold": old_model_queue_threshold,
"oldModelSlotsBeforeScan": old_model_slots_before_scan,
"olderHistoryScanEnabled": allow_older_models_for_scan,
},
"scanLimit": scan_limit,
"candidateGoal": candidate_goal,
"scanStages": scan_stages,

View File

@@ -9,7 +9,7 @@ from datetime import datetime, timedelta
from pathlib import Path
from typing import Any
from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, write_json
from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, utc_now, write_json
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
from http_json import HttpJsonError, JsonHttpClient
@@ -21,6 +21,10 @@ class ModelHubAPIError(RuntimeError):
self.payload = payload
class OldModelQueuePolicyError(ModelHubAPIError):
"""No account may accept an older model under the per-account queue policy."""
def parse_model_submission_precheck(payload: Any) -> dict[str, Any]:
"""Validate the community lookup response instead of failing open."""
if not isinstance(payload, dict):
@@ -385,12 +389,26 @@ class ModelHubClientPool:
capacity_probe_interval_cycles: int = 3,
capacity_probe_cooldown_cycles: int = 3,
capacity_state_path: Path | str | None = None,
old_model_queue_threshold: int | None = None,
recent_model_days: int | None = None,
) -> None:
if not clients:
raise ValueError("At least one ModelHub client is required")
self.clients = clients
configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100")
self.active_task_cap = max(1, int(configured_cap))
configured_old_threshold = (
old_model_queue_threshold
if old_model_queue_threshold is not None
else os.getenv("MODELHUB_OLD_MODEL_QUEUE_THRESHOLD", "80")
)
configured_recent_days = (
recent_model_days
if recent_model_days is not None
else os.getenv("MODELHUB_RECENT_MODEL_DAYS", "7")
)
self.old_model_queue_threshold = max(1, int(configured_old_threshold))
self.recent_model_days = max(1, int(configured_recent_days))
self._capacity_state_path = Path(capacity_state_path) if capacity_state_path else None
self._account_keys = [self._account_key(client, index) for index, client in enumerate(clients)]
self._account_caps = self._load_account_caps(self.active_task_cap)
@@ -671,12 +689,21 @@ class ModelHubClientPool:
def processed_gpus_for_model(self, model_id: str) -> set[str]:
return set(self.model_submission_precheck(model_id)["processedGpus"])
def _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None:
def _reserve_account(
self,
excluded: set[int],
*,
effective_count_below: int | None = None,
) -> tuple[int, int] | None:
with self._state_lock:
remaining_by_index = {
index: self._account_caps[index] - self._effective_count_locked(index)
for index in range(len(self.clients))
if index not in excluded
and (
effective_count_below is None
or self._effective_count_locked(index) < effective_count_below
)
}
usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
if not usable:
@@ -696,7 +723,12 @@ class ModelHubClientPool:
}
return selected_index, reservation_id
def _reserve_probe_account(self, excluded: set[int]) -> tuple[int, int] | None:
def _reserve_probe_account(
self,
excluded: set[int],
*,
effective_count_below: int | None = None,
) -> tuple[int, int] | None:
with self._state_lock:
if not self._capacity_probe_enabled:
return None
@@ -707,6 +739,10 @@ class ModelHubClientPool:
and index not in self._capacity_probe_attempted
and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
and self._effective_count_locked(index) >= self._account_caps[index]
and (
effective_count_below is None
or self._effective_count_locked(index) < effective_count_below
)
]
if not eligible:
return None
@@ -763,17 +799,58 @@ class ModelHubClientPool:
self._persist_account_caps()
print(f"[capacity] account={index + 1:02d} discovered_cap={discovered_cap}", flush=True)
def old_model_submit_slots(self) -> int:
"""Return queue positions where models older than the recent window are allowed."""
self._refresh_active_counts()
with self._state_lock:
return sum(
max(
0,
min(self._account_caps[index], self.old_model_queue_threshold)
- self._effective_count_locked(index),
)
for index in range(len(self.clients))
)
def can_submit_old_models(self) -> bool:
return self.old_model_submit_slots() > 0
def add_task_for_model(
self,
payload: dict[str, Any],
*,
model_last_modified: datetime | str | None,
submitted_at: datetime | str | None = None,
) -> dict[str, Any]:
reference_time = parse_datetime(submitted_at) or utc_now()
last_modified = parse_datetime(model_last_modified)
is_old_model = (
last_modified is None
or last_modified < reference_time - timedelta(days=self.recent_model_days)
)
return self._add_task(payload, old_model_only=is_old_model)
def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
return self._add_task(payload, old_model_only=False)
def _add_task(self, payload: dict[str, Any], *, old_model_only: bool) -> dict[str, Any]:
self._refresh_active_counts()
attempted_accounts: set[int] = set()
forced_refresh_done = False
last_capacity_error: ModelHubAPIError | None = None
account_count_limit = self.old_model_queue_threshold if old_model_only else None
while True:
reservation = self._reserve_account(attempted_accounts)
reservation = self._reserve_account(
attempted_accounts,
effective_count_below=account_count_limit,
)
capacity_probe = False
if reservation is None:
reservation = self._reserve_probe_account(attempted_accounts)
reservation = self._reserve_probe_account(
attempted_accounts,
effective_count_below=account_count_limit,
)
capacity_probe = reservation is not None
if reservation is None:
if not forced_refresh_done:
@@ -782,6 +859,11 @@ class ModelHubClientPool:
continue
if last_capacity_error is not None:
raise last_capacity_error
if old_model_only:
raise OldModelQueuePolicyError(
f"所有账号活动队列均已达到 {self.old_model_queue_threshold}"
f"仅允许提交最近 {self.recent_model_days} 天内更新的模型"
)
raise ModelHubAPIError(
"当前等待中或运行中的异步模型验证任务数量已达已知上限"
)

View File

@@ -58,6 +58,24 @@ def build_parser() -> argparse.ArgumentParser:
default=0,
help="Maximum tasks to submit in one cycle (0 means unlimited)",
)
parser.add_argument(
"--old-model-queue-threshold",
type=int,
default=int(os.getenv("MODELHUB_OLD_MODEL_QUEUE_THRESHOLD", "80")),
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-threshold",
type=int,
default=int(os.getenv("MODELHUB_DYNAMIC_OLD_MODEL_CLEANUP_THRESHOLD", "95")),
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)
@@ -206,9 +224,30 @@ def _build_modelhub_client(base_args: argparse.Namespace) -> 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)),
old_model_queue_threshold=max(1, int(getattr(base_args, "old_model_queue_threshold", 80) or 80)),
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]:
"""Use the strict admission boundary once, then retain a 15-slot buffer."""
initial_threshold = max(
1,
int(getattr(base_args, "old_model_queue_threshold", 80) or 80),
)
dynamic_threshold = max(
initial_threshold,
int(getattr(base_args, "dynamic_old_model_cleanup_threshold", 95) or 95),
)
if initial_cleanup_pending:
return "initial", initial_threshold
return "dynamic", dynamic_threshold
def run_poll_loop(
*,
base_args: argparse.Namespace,
@@ -243,6 +282,7 @@ def run_poll_loop(
submitted_total = 0
cycles = 0
stopped_reason = "max_cycles_reached"
initial_age_cleanup_pending = True
while True:
if base_args.max_cycles and cycles >= base_args.max_cycles:
@@ -263,12 +303,22 @@ def run_poll_loop(
try:
cleanup_feedback = outcome_tracker.get_stats_report()
cleanup_gpu_memory = cleanup_feedback.get("observedGpuMemoryGiB") or {}
age_cleanup_mode, age_cleanup_threshold = resolve_age_cleanup_policy(
base_args,
initial_cleanup_pending=initial_age_cleanup_pending,
)
log(
f"[queue-cleanup] mode={age_cleanup_mode} "
f"old_model_threshold={age_cleanup_threshold} "
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,
age_queue_threshold=age_cleanup_threshold,
log=log,
)
write_json(
@@ -284,12 +334,16 @@ def run_poll_loop(
queue_cleanup_runs.append(
{
"cycle": cycles,
"mode": age_cleanup_mode,
"ageQueueThreshold": age_cleanup_threshold,
"activeScanned": cleanup_summary["activeScanned"],
"certainOomCount": cleanup_summary["certainOomCount"],
"oldOverflowCount": cleanup_summary["oldOverflowCount"],
"cancelledCount": cleanup_summary["cancelledCount"],
"stopErrorCount": len(cleanup_summary["stopErrors"]),
}
)
initial_age_cleanup_pending = False
except Exception as exc:
log(f"[queue-cleanup] error={type(exc).__name__}: {exc} continue_polling=true")

View File

@@ -3,11 +3,12 @@ from __future__ import annotations
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Callable, Iterable
from candidate_preflight import CandidatePreflightAdvisor, MODEL_LOAD_OVERHEAD
from common import write_json
from common import utc_now, write_json
from hf_discovery import HuggingFaceDiscovery, inspect_repo_tree
from modelhub_client import ModelHubClient, ModelHubClientPool
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
@@ -159,6 +160,41 @@ def _load_repository_sizes(
return sizes, errors
def _load_model_last_modified(
model_ids: set[str],
*,
discovery: HuggingFaceDiscovery,
read_concurrency: int,
log: Callable[[str], None],
) -> tuple[dict[str, datetime], dict[str, str]]:
values: dict[str, datetime] = {}
errors: dict[str, str] = {}
completed = 0
workers = min(max(1, int(read_concurrency)), max(1, len(model_ids)))
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(discovery.get_model_last_modified, model_id): model_id
for model_id in sorted(model_ids)
}
for future in as_completed(futures):
model_id = futures[future]
try:
value = future.result()
if value is None:
errors[model_id] = "model_last_modified_unknown"
else:
values[model_id] = value
except Exception as exc:
errors[model_id] = f"{type(exc).__name__}: {exc}"
completed += 1
if completed == len(model_ids) or completed % 50 == 0:
log(
f"[queue-cleanup] age_scan={completed}/{len(model_ids)} "
f"complete={len(values)} unknown={len(errors)}"
)
return values, errors
def find_certain_oom_tasks(
tasks: list[OwnedTask],
*,
@@ -202,6 +238,62 @@ def find_certain_oom_tasks(
return decisions, skipped
def find_old_overflow_tasks(
tasks: list[OwnedTask],
*,
model_last_modified: dict[str, datetime],
queue_threshold: int = 80,
recent_model_days: int = 7,
reference_time: datetime | None = None,
incomplete_accounts: set[int] | None = None,
) -> tuple[list[dict[str, Any]], dict[str, int]]:
threshold = max(1, int(queue_threshold))
recent_days = max(1, int(recent_model_days))
cutoff = (reference_time or utc_now()) - timedelta(days=recent_days)
incomplete_accounts = incomplete_accounts or set()
grouped: dict[int, list[OwnedTask]] = {}
for task in tasks:
grouped.setdefault(task.account_index, []).append(task)
decisions: list[dict[str, Any]] = []
skipped = {
"accountsWithIncompleteListing": len(incomplete_accounts),
"withinFirstQueuePositions": 0,
"recentOverflowTasks": 0,
"modelAgeUnknown": 0,
}
for account_index, account_tasks in sorted(grouped.items()):
if account_index in incomplete_accounts:
continue
ordered = sorted(account_tasks, key=lambda item: item.task_id)
skipped["withinFirstQueuePositions"] += min(threshold, len(ordered))
for position, task in enumerate(ordered, start=1):
if position <= threshold:
continue
last_modified = model_last_modified.get(task.model_id)
if last_modified is None:
skipped["modelAgeUnknown"] += 1
continue
if last_modified >= cutoff:
skipped["recentOverflowTasks"] += 1
continue
decisions.append(
{
"accountIndex": account_index + 1,
"taskId": task.task_id,
"modelId": task.model_id,
"gpuType": task.gpu_type,
"status": task.status,
"queuePosition": position,
"queueThreshold": threshold,
"modelLastModified": last_modified.isoformat(),
"recentCutoff": cutoff.isoformat(),
"reason": "old_model_beyond_account_queue_threshold",
}
)
return decisions, skipped
def _chunks(values: list[int], size: int) -> Iterable[list[int]]:
for offset in range(0, len(values), size):
yield values[offset : offset + size]
@@ -215,10 +307,20 @@ def cleanup_certain_oom_tasks(
read_concurrency: int = 6,
stop_batch_size: int = DEFAULT_STOP_BATCH_SIZE,
gpu_memory_gib: dict[str, float] | None = None,
age_queue_threshold: int | None = None,
reference_time: datetime | None = None,
log: Callable[[str], None] = print,
) -> dict[str, Any]:
"""Stop only active tasks that cannot fit the selected GPU by known capacity."""
"""Stop deterministic OOM tasks and old tasks beyond each account's protected prefix."""
clients = list(modelhub.clients)
reference_time = reference_time or utc_now()
configured_queue_threshold = (
age_queue_threshold
if age_queue_threshold is not None
else getattr(modelhub, "old_model_queue_threshold", 80)
)
queue_threshold = max(1, int(configured_queue_threshold or 80))
recent_model_days = max(1, int(getattr(modelhub, "recent_model_days", 7) or 7))
tasks, listing_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
log(
f"[queue-cleanup] active_scanned={len(tasks)} accounts={len(clients)} "
@@ -232,27 +334,104 @@ def cleanup_certain_oom_tasks(
read_concurrency=read_concurrency,
log=log,
)
decisions, skipped = find_certain_oom_tasks(
oom_decisions, skipped = find_certain_oom_tasks(
tasks,
repository_sizes=repository_sizes,
gpu_memory_gib=gpu_memory_gib,
)
log(
f"[queue-cleanup] certain_oom={len(decisions)} "
f"[queue-cleanup] certain_oom={len(oom_decisions)} "
f"fits={skipped['fitsKnownCapacity']} size_unknown={skipped['repositorySizeUnknown']} "
f"gpu_unknown={skipped['gpuCapacityUnknown']} dry_run={str(bool(dry_run)).lower()}"
)
oom_task_keys = {
(int(decision["accountIndex"]) - 1, int(decision["taskId"]))
for decision in oom_decisions
}
# OOM tasks are stopped first. Rank the age-policy queue as it will look
# after those certain failures are gone, so an old task moving into the
# protected first N positions is not over-cancelled.
age_rank_tasks = [
task for task in tasks if (task.account_index, task.task_id) not in oom_task_keys
]
overflow_model_ids = {
task.model_id
for account_index in range(len(clients))
if account_index not in listing_errors
for task in sorted(
(item for item in age_rank_tasks if item.account_index == account_index),
key=lambda item: item.task_id,
)[queue_threshold:]
}
model_last_modified, age_errors = _load_model_last_modified(
overflow_model_ids,
discovery=discovery,
read_concurrency=read_concurrency,
log=log,
)
old_overflow_decisions, age_skipped = find_old_overflow_tasks(
age_rank_tasks,
model_last_modified=model_last_modified,
queue_threshold=queue_threshold,
recent_model_days=recent_model_days,
reference_time=reference_time,
incomplete_accounts=set(listing_errors),
)
log(
f"[queue-cleanup] old_overflow={len(old_overflow_decisions)} "
f"threshold={queue_threshold} recent_days={recent_model_days} "
f"recent_overflow={age_skipped['recentOverflowTasks']} "
f"age_unknown={age_skipped['modelAgeUnknown']}"
)
decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {}
for decision in oom_decisions:
enriched = dict(decision)
enriched["cleanupReasons"] = [decision["reason"]]
decisions_by_key[(int(decision["accountIndex"]), int(decision["taskId"]))] = enriched
for decision in old_overflow_decisions:
key = (int(decision["accountIndex"]), int(decision["taskId"]))
existing = decisions_by_key.get(key)
if existing is None:
enriched = dict(decision)
enriched["cleanupReasons"] = [decision["reason"]]
decisions_by_key[key] = enriched
continue
existing["cleanupReasons"].append(decision["reason"])
for field in ("queuePosition", "queueThreshold", "modelLastModified", "recentCutoff"):
existing[field] = decision[field]
decisions = sorted(
decisions_by_key.values(),
key=lambda item: (int(item["accountIndex"]), int(item["taskId"])),
)
cancelled: list[dict[str, Any]] = []
disappeared: list[dict[str, Any]] = []
policy_no_longer_applies: list[dict[str, Any]] = []
stop_errors: list[dict[str, Any]] = []
if decisions and not dry_run:
# Re-read each account immediately before mutation. If any active-state
# query fails for that account, fail closed and do not terminate its tasks.
refreshed_tasks, refresh_errors = collect_active_tasks(clients, read_concurrency=read_concurrency)
active_ids_by_account: dict[int, set[int]] = {}
active_positions_by_account: dict[int, dict[int, int]] = {}
for task in refreshed_tasks:
active_ids_by_account.setdefault(task.account_index, set()).add(task.task_id)
for account_index in range(len(clients)):
planned_oom_ids = {
int(decision["taskId"])
for decision in oom_decisions
if int(decision["accountIndex"]) - 1 == account_index
}
ordered_ids = sorted(
task.task_id
for task in refreshed_tasks
if task.account_index == account_index and task.task_id not in planned_oom_ids
)
active_positions_by_account[account_index] = {
task_id: position for position, task_id in enumerate(ordered_ids, start=1)
}
for account_index, details in sorted(refresh_errors.items()):
stop_errors.append(
@@ -271,6 +450,16 @@ def cleanup_certain_oom_tasks(
if int(decision["taskId"]) not in active_ids_by_account.get(account_index, set()):
disappeared.append(decision)
continue
cleanup_reasons = set(decision.get("cleanupReasons") or [decision.get("reason")])
age_only = cleanup_reasons == {"old_model_beyond_account_queue_threshold"}
current_position = active_positions_by_account.get(account_index, {}).get(int(decision["taskId"]))
if age_only and (current_position is None or current_position <= queue_threshold):
policy_no_longer_applies.append(
{**decision, "recheckedQueuePosition": current_position}
)
continue
if current_position is not None:
decision["recheckedQueuePosition"] = current_position
by_account.setdefault(account_index, []).append(decision)
batch_size = max(1, min(100, int(stop_batch_size)))
@@ -279,25 +468,36 @@ def cleanup_certain_oom_tasks(
if stop_failed:
break
decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]}
task_ids = sorted(decisions_by_id)
for batch in _chunks(task_ids, batch_size):
try:
clients[account_index].stop_tasks(batch)
except Exception as exc:
stop_errors.append(
{
"accountIndex": account_index + 1,
"taskIds": batch,
"error": f"{type(exc).__name__}: {exc}",
}
)
stop_failed = True
break
cancelled.extend(decisions_by_id[task_id] for task_id in batch)
log(
f"[queue-cleanup] account={account_index + 1:02d} "
f"cancelled_batch={len(batch)} cancelled_total={len(cancelled)}"
phase_ids = [
sorted(
task_id
for task_id, decision in decisions_by_id.items()
if ("certain_oom_repository_size_exceeds_gpu_capacity" in decision["cleanupReasons"])
== is_oom_phase
)
for is_oom_phase in (True, False)
]
for task_ids in phase_ids:
for batch in _chunks(task_ids, batch_size):
try:
clients[account_index].stop_tasks(batch)
except Exception as exc:
stop_errors.append(
{
"accountIndex": account_index + 1,
"taskIds": batch,
"error": f"{type(exc).__name__}: {exc}",
}
)
stop_failed = True
break
cancelled.extend(decisions_by_id[task_id] for task_id in batch)
log(
f"[queue-cleanup] account={account_index + 1:02d} "
f"cancelled_batch={len(batch)} cancelled_total={len(cancelled)}"
)
if stop_failed:
break
if cancelled and hasattr(modelhub, "refresh_active_counts"):
modelhub.refresh_active_counts()
@@ -309,19 +509,29 @@ def cleanup_certain_oom_tasks(
"uniqueModels": len(model_ids),
"repositorySizesComplete": len(repository_sizes),
"repositorySizeErrors": size_errors,
"modelAgeMetadataComplete": len(model_last_modified),
"modelAgeErrors": age_errors,
"listingErrors": {str(index + 1): values for index, values in listing_errors.items()},
"certainOomCount": len(decisions),
"certainOomTasks": decisions,
"certainOomCount": len(oom_decisions),
"certainOomTasks": oom_decisions,
"oldOverflowCount": len(old_overflow_decisions),
"oldOverflowTasks": old_overflow_decisions,
"oldModelQueueThreshold": queue_threshold,
"recentModelDays": recent_model_days,
"agePolicySkipped": age_skipped,
"cleanupCandidateCount": len(decisions),
"skipped": skipped,
"cancelledCount": len(cancelled),
"cancelledTasks": cancelled,
"noLongerActiveCount": len(disappeared),
"policyNoLongerAppliesCount": len(policy_no_longer_applies),
"policyNoLongerAppliesTasks": policy_no_longer_applies,
"stopErrors": stop_errors,
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Safely stop active ModelHub tasks that are certain to run out of GPU memory.")
parser = argparse.ArgumentParser(description="Safely stop deterministic OOM and over-threshold old ModelHub tasks.")
parser.add_argument("--execute", action="store_true", help="Actually terminate selected tasks; otherwise only print a preview")
parser.add_argument("--read-concurrency", type=int, default=6)
parser.add_argument("--stop-batch-size", type=int, default=DEFAULT_STOP_BATCH_SIZE)
@@ -353,6 +563,7 @@ def main(argv: list[str] | None = None) -> int:
write_json(report_path, summary)
print(
f"[queue-cleanup] finished certain_oom={summary['certainOomCount']} "
f"old_overflow={summary['oldOverflowCount']} "
f"cancelled={summary['cancelledCount']} stop_errors={len(summary['stopErrors'])} "
f"report={report_path}",
flush=True,

View File

@@ -1 +1 @@
AGENT_VERSION = "2026.08.11.2"
AGENT_VERSION = "2026.08.11.3"

View File

@@ -17,8 +17,13 @@ if str(PACKAGE_DIR) in sys.path:
sys.path.insert(0, str(PACKAGE_DIR))
from common import read_json, read_jsonl # noqa: E402
from main import build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402
from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_submission_error # 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
@@ -210,6 +215,109 @@ def make_candidate(index: int) -> dict:
class ClientPoolConcurrencyTests(unittest.TestCase):
def test_old_models_use_only_each_accounts_first_eighty_queue_positions(self) -> None:
below_threshold = FakeClient(active_count=79)
at_threshold = FakeClient(active_count=80)
pool = ModelHubClientPool(
[below_threshold, at_threshold], # type: ignore[list-item]
active_task_cap=100,
active_counts_ttl=60,
old_model_queue_threshold=80,
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-80"},
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_cross_eighty(self) -> None:
client = FakeClient(active_count=78)
pool = ModelHubClientPool(
[client], # type: ignore[list-item]
active_task_cap=100,
active_counts_ttl=60,
old_model_queue_threshold=80,
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_eighty(self) -> None:
client = FakeClient(active_count=80)
pool = ModelHubClientPool(
[client], # type: ignore[list-item]
active_task_cap=100,
old_model_queue_threshold=80,
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("old_model_policy_skipped", result["outcome"])
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)

View File

@@ -47,7 +47,27 @@ class RateLimitedPageClient:
return {"success": True, "data": {"models": items}}
class ModelMetadataClient:
def __init__(self) -> None:
self.calls = 0
def request_json(self, _method: str, _path: str, **_kwargs) -> dict: # noqa: ANN003
self.calls += 1
return {"Code": 200, "Data": {"LastUpdatedTime": 1786294389}}
class ModelScopeDiscoveryTests(unittest.TestCase):
def test_model_last_modified_uses_detail_api_and_is_cached(self) -> None:
metadata_client = ModelMetadataClient()
discovery = HuggingFaceDiscovery(legacy_http_client=metadata_client) # type: ignore[arg-type]
first = discovery.get_model_last_modified("owner/model")
second = discovery.get_model_last_modified("owner/model")
self.assertEqual(first, second)
self.assertEqual(1786294389, int(first.timestamp())) # type: ignore[union-attr]
self.assertEqual(1, metadata_client.calls)
def test_openapi_page_size_never_exceeds_platform_limit(self) -> None:
http_client = RecordingHttpClient()
discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type]

46
tests/test_poll_policy.py Normal file
View File

@@ -0,0 +1,46 @@
from __future__ import annotations
import argparse
import sys
import unittest
from pathlib import Path
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 poll_runner import resolve_age_cleanup_policy # noqa: E402
class PollPolicyTests(unittest.TestCase):
def test_age_cleanup_uses_eighty_once_then_ninety_five(self) -> None:
args = argparse.Namespace(
old_model_queue_threshold=80,
dynamic_old_model_cleanup_threshold=95,
)
self.assertEqual(
("initial", 80),
resolve_age_cleanup_policy(args, initial_cleanup_pending=True),
)
self.assertEqual(
("dynamic", 95),
resolve_age_cleanup_policy(args, initial_cleanup_pending=False),
)
def test_dynamic_cleanup_cannot_be_stricter_than_admission(self) -> None:
args = argparse.Namespace(
old_model_queue_threshold=80,
dynamic_old_model_cleanup_threshold=70,
)
self.assertEqual(
("dynamic", 80),
resolve_age_cleanup_policy(args, initial_cleanup_pending=False),
)
if __name__ == "__main__":
unittest.main()

View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import unittest
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
@@ -12,16 +13,28 @@ if str(PACKAGE_DIR) in sys.path:
sys.path.insert(0, str(PACKAGE_DIR))
from modelhub_client import ModelHubClient, ModelHubClientPool # noqa: E402
from queue_cleanup import OwnedTask, cleanup_certain_oom_tasks, find_certain_oom_tasks # noqa: E402
from queue_cleanup import ( # noqa: E402
OwnedTask,
cleanup_certain_oom_tasks,
find_certain_oom_tasks,
find_old_overflow_tasks,
)
GIB = 1024**3
class FakeQueueClient:
def __init__(self, records: list[dict[str, Any]], *, disappear_on_recheck: bool = False) -> None:
def __init__(
self,
records: list[dict[str, Any]],
*,
disappear_on_recheck: bool = False,
drop_first_on_recheck: bool = False,
) -> None:
self.records = list(records)
self.disappear_on_recheck = disappear_on_recheck
self.drop_first_on_recheck = drop_first_on_recheck
self.waiting_reads = 0
self.stopped: list[list[int]] = []
@@ -32,6 +45,8 @@ class FakeQueueClient:
records: list[dict[str, Any]] = []
else:
records = [record for record in self.records if record["status"] == "waiting"]
if self.drop_first_on_recheck and self.waiting_reads >= 2:
records = sorted(records, key=lambda item: int(item["taskId"]))[1:]
else:
records = [record for record in self.records if record["status"] == status]
return {"code": 0, "data": {"records": records, "pages": 1}}
@@ -47,8 +62,13 @@ class FakeQueueClient:
class FakeDiscovery:
def __init__(self, sizes: dict[str, int | None]) -> None:
def __init__(
self,
sizes: dict[str, int | None],
last_modified: dict[str, datetime | None] | None = None,
) -> None:
self.sizes = sizes
self.last_modified = last_modified or {}
def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]:
size = self.sizes[repo_id]
@@ -56,6 +76,9 @@ class FakeDiscovery:
return [{"Path": "model.safetensors"}]
return [{"Path": "model.safetensors", "Size": size}]
def get_model_last_modified(self, repo_id: str) -> datetime | None:
return self.last_modified.get(repo_id)
class RecordingHttpClient:
def __init__(self) -> None:
@@ -74,6 +97,151 @@ class RecordingHttpClient:
class QueueCleanupTests(unittest.TestCase):
def test_old_models_are_selected_only_after_each_accounts_first_eighty_tasks(self) -> None:
now = datetime(2026, 8, 11, tzinfo=timezone.utc)
tasks = [
OwnedTask(0, index, "owner/old", "Iluvatar_bi-100", "waiting")
for index in range(1, 82)
]
tasks.append(OwnedTask(0, 82, "owner/recent", "Iluvatar_bi-100", "waiting"))
selected, skipped = find_old_overflow_tasks(
tasks,
model_last_modified={
"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc),
"owner/recent": datetime(2026, 8, 10, tzinfo=timezone.utc),
},
queue_threshold=80,
recent_model_days=7,
reference_time=now,
)
self.assertEqual([81], [item["taskId"] for item in selected])
self.assertEqual(81, selected[0]["queuePosition"])
self.assertEqual(1, skipped["recentOverflowTasks"])
def test_old_overflow_task_is_not_stopped_if_it_moves_into_first_eighty(self) -> None:
records = [
{
"taskId": index,
"modelId": "owner/old",
"gpuType": "Iluvatar_bi-100",
"status": "waiting",
}
for index in range(1, 82)
]
client = FakeQueueClient(records, drop_first_on_recheck=True)
pool = ModelHubClientPool(
[client], # type: ignore[list-item]
active_task_cap=100,
old_model_queue_threshold=80,
)
summary = cleanup_certain_oom_tasks(
pool,
FakeDiscovery(
{"owner/old": 1 * GIB},
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
), # type: ignore[arg-type]
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
log=lambda _message: None,
)
self.assertEqual(1, summary["oldOverflowCount"])
self.assertEqual(0, summary["cancelledCount"])
self.assertEqual(1, summary["policyNoLongerAppliesCount"])
self.assertEqual([], client.stopped)
def test_cleanup_stops_old_task_beyond_eightieth_position(self) -> None:
records = [
{
"taskId": index,
"modelId": "owner/old",
"gpuType": "Iluvatar_bi-100",
"status": "waiting",
}
for index in range(1, 82)
]
client = FakeQueueClient(records)
pool = ModelHubClientPool(
[client], # type: ignore[list-item]
active_task_cap=100,
old_model_queue_threshold=80,
)
summary = cleanup_certain_oom_tasks(
pool,
FakeDiscovery(
{"owner/old": 1 * GIB},
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
), # type: ignore[arg-type]
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
log=lambda _message: None,
)
self.assertEqual(1, summary["oldOverflowCount"])
self.assertEqual(1, summary["cancelledCount"])
self.assertEqual([[81]], client.stopped)
def test_dynamic_cleanup_keeps_positions_through_ninety_five(self) -> None:
records = [
{
"taskId": index,
"modelId": "owner/old",
"gpuType": "Iluvatar_bi-100",
"status": "waiting",
}
for index in range(1, 97)
]
client = FakeQueueClient(records)
pool = ModelHubClientPool(
[client], # type: ignore[list-item]
active_task_cap=100,
old_model_queue_threshold=80,
)
summary = cleanup_certain_oom_tasks(
pool,
FakeDiscovery(
{"owner/old": 1 * GIB},
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
), # type: ignore[arg-type]
age_queue_threshold=95,
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
log=lambda _message: None,
)
self.assertEqual(95, summary["oldModelQueueThreshold"])
self.assertEqual([96], [item["taskId"] for item in summary["oldOverflowTasks"]])
self.assertEqual([[96]], client.stopped)
def test_oom_is_removed_before_recalculating_old_overflow_positions(self) -> None:
records = [
{
"taskId": index,
"modelId": "owner/large" if index == 1 else "owner/old",
"gpuType": "Iluvatar_bi-100",
"status": "waiting",
}
for index in range(1, 83)
]
client = FakeQueueClient(records)
pool = ModelHubClientPool(
[client], # type: ignore[list-item]
active_task_cap=100,
old_model_queue_threshold=80,
)
summary = cleanup_certain_oom_tasks(
pool,
FakeDiscovery(
{"owner/large": 40 * GIB, "owner/old": 1 * GIB},
{"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
), # type: ignore[arg-type]
reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
log=lambda _message: None,
)
self.assertEqual(1, summary["certainOomCount"])
self.assertEqual([82], [item["taskId"] for item in summary["oldOverflowTasks"]])
self.assertEqual([[1], [82]], client.stopped)
def test_stop_tasks_uses_documented_put_endpoint_and_integer_ids(self) -> None:
http = RecordingHttpClient()
client = ModelHubClient(http_client=http) # type: ignore[arg-type]