feat: reserve queue capacity for recent models
This commit is contained in:
30
README.md
30
README.md
@@ -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 81–100. 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
|
||||
```
|
||||
|
||||
@@ -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 1–80 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 81–95 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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
"当前等待中或运行中的异步模型验证任务数量已达已知上限"
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -1 +1 @@
|
||||
AGENT_VERSION = "2026.08.11.2"
|
||||
AGENT_VERSION = "2026.08.11.3"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
46
tests/test_poll_policy.py
Normal 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()
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user