feat: add tiered durable state and memory bounds
This commit is contained in:
@@ -79,6 +79,12 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
default=4,
|
||||
help="Multiplier used when auto-deriving scan limit from quota/queue capacity",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--candidate-pool-limit",
|
||||
type=int,
|
||||
default=100,
|
||||
help="Maximum prepared candidates retained before submitting; bounds memory in large queues",
|
||||
)
|
||||
parser.add_argument("--min-downloads", type=int, default=50, help="Minimum downloads threshold")
|
||||
parser.add_argument("--daily-target", type=int, default=0, help="Daily submission target across all auto runs; 0 means unlimited")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Only write artifacts without creating tasks")
|
||||
@@ -746,6 +752,7 @@ def collect_candidates_from_models(
|
||||
allow_dynamic_tasks: bool = False,
|
||||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]], int]:
|
||||
candidates: list[dict[str, Any]] = []
|
||||
candidate_models: set[str] = set()
|
||||
skipped: list[dict[str, Any]] = []
|
||||
failed: list[dict[str, Any]] = []
|
||||
processed_count = 0
|
||||
@@ -754,7 +761,7 @@ def collect_candidates_from_models(
|
||||
new_models = [model for model in models if model.repo_id not in seen_model_ids]
|
||||
|
||||
for offset in range(0, len(new_models), chunk_size):
|
||||
if len(candidates) >= candidate_goal:
|
||||
if len(candidate_models) >= candidate_goal:
|
||||
break
|
||||
chunk = new_models[offset : offset + chunk_size]
|
||||
for model in chunk:
|
||||
@@ -794,6 +801,10 @@ def collect_candidates_from_models(
|
||||
for index in range(len(chunk)):
|
||||
model_candidates, model_skipped, model_failed = ordered_results.get(index, ([], [], []))
|
||||
candidates.extend(model_candidates)
|
||||
if model_candidates:
|
||||
model_id = str(model_candidates[0].get("repoId") or model_candidates[0].get("modelAddress") or "")
|
||||
if model_id:
|
||||
candidate_models.add(model_id)
|
||||
skipped.extend(model_skipped)
|
||||
failed.extend(model_failed)
|
||||
processed_count += len(chunk)
|
||||
@@ -927,6 +938,46 @@ def one_candidate_per_model(candidates: list[dict[str, Any]]) -> list[dict[str,
|
||||
return selected
|
||||
|
||||
|
||||
def limit_routes_per_model(
|
||||
candidates: list[dict[str, Any]],
|
||||
*,
|
||||
max_routes: int = 3,
|
||||
) -> list[dict[str, Any]]:
|
||||
selected: list[dict[str, Any]] = []
|
||||
counts: Counter[str] = Counter()
|
||||
for candidate in candidates:
|
||||
model_id = str(candidate.get("repoId") or candidate.get("modelAddress") or "")
|
||||
if not model_id or counts[model_id] >= max(1, int(max_routes)):
|
||||
continue
|
||||
counts[model_id] += 1
|
||||
selected.append(candidate)
|
||||
return selected
|
||||
|
||||
|
||||
def candidate_model_count(candidates: list[dict[str, Any]]) -> int:
|
||||
return len(
|
||||
{
|
||||
str(candidate.get("repoId") or candidate.get("modelAddress") or "")
|
||||
for candidate in candidates
|
||||
if candidate.get("repoId") or candidate.get("modelAddress")
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def limit_candidate_models(candidates: list[dict[str, Any]], *, max_models: int) -> list[dict[str, Any]]:
|
||||
selected: list[dict[str, Any]] = []
|
||||
models: set[str] = set()
|
||||
for candidate in candidates:
|
||||
model_id = str(candidate.get("repoId") or candidate.get("modelAddress") or "")
|
||||
if not model_id:
|
||||
continue
|
||||
if model_id not in models and len(models) >= max(0, int(max_models)):
|
||||
continue
|
||||
models.add(model_id)
|
||||
selected.append(candidate)
|
||||
return selected
|
||||
|
||||
|
||||
def run_submission(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
@@ -1271,7 +1322,14 @@ def run_submission(
|
||||
if strategy_manager is not None:
|
||||
submission_goal = min(submission_goal, strategy_manager.submissions_until_refresh)
|
||||
attempt_multiplier = max(1, int(getattr(args, "scan_multiplier", 4) or 1))
|
||||
candidate_goal = max(submission_goal, submission_goal * attempt_multiplier)
|
||||
raw_candidate_goal = max(submission_goal, submission_goal * attempt_multiplier)
|
||||
candidate_pool_limit = max(20, int(getattr(args, "candidate_pool_limit", 100) or 100))
|
||||
candidate_goal = min(raw_candidate_goal, candidate_pool_limit)
|
||||
print(
|
||||
f"[scan] candidate_pool_limit={candidate_pool_limit} "
|
||||
f"raw_goal={raw_candidate_goal} effective_goal={candidate_goal}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
pipeline_tags = pipeline_tags_for_task_types(selected_task_types)
|
||||
if official_registry is not None and not getattr(args, "task_types", None):
|
||||
@@ -1305,7 +1363,8 @@ def run_submission(
|
||||
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:
|
||||
current_candidate_models = candidate_model_count(candidates)
|
||||
if candidate_goal <= 0 or current_candidate_models >= candidate_goal:
|
||||
break
|
||||
skipped_before_stage = len(skipped)
|
||||
stage_updated_after = stage["updatedAfter"]
|
||||
@@ -1348,7 +1407,7 @@ def run_submission(
|
||||
stage_candidates, stage_skipped, stage_failed, processed_count = collect_candidates_from_models(
|
||||
models=models,
|
||||
seen_model_ids=seen_model_ids,
|
||||
candidate_goal=max(1, candidate_goal - len(candidates)),
|
||||
candidate_goal=max(1, candidate_goal - current_candidate_models),
|
||||
hf_discovery=hf_discovery,
|
||||
modelhub_client=modelhub_client,
|
||||
template_selector=template_selector,
|
||||
@@ -1363,6 +1422,12 @@ def run_submission(
|
||||
config_optimizer=config_optimizer,
|
||||
allow_dynamic_tasks=not bool(getattr(args, "task_types", None)),
|
||||
)
|
||||
stage_candidates = limit_routes_per_model(stage_candidates, max_routes=3)
|
||||
remaining_candidate_capacity = max(0, candidate_goal - current_candidate_models)
|
||||
stage_candidates = limit_candidate_models(
|
||||
stage_candidates,
|
||||
max_models=remaining_candidate_capacity,
|
||||
)
|
||||
candidates.extend(stage_candidates)
|
||||
skipped.extend(stage_skipped)
|
||||
failed.extend(stage_failed)
|
||||
@@ -1375,13 +1440,15 @@ def run_submission(
|
||||
"newModelsProcessed": processed_count,
|
||||
"candidatesAdded": len(stage_candidates),
|
||||
"candidateCountAfterStage": len(candidates),
|
||||
"candidateModelsAfterStage": candidate_model_count(candidates),
|
||||
"skippedAdded": len(skipped) - skipped_before_stage,
|
||||
"failedAdded": len(stage_failed),
|
||||
}
|
||||
scan_stages.append(stage_summary)
|
||||
print(
|
||||
f"[scan] stage={stage['name']} discovered={len(models)} new_processed={processed_count} "
|
||||
f"candidates_added={len(stage_candidates)} candidates_total={len(candidates)}/{candidate_goal} "
|
||||
f"routes_added={len(stage_candidates)} routes_total={len(candidates)} "
|
||||
f"candidate_models={candidate_model_count(candidates)}/{candidate_goal} "
|
||||
f"skipped_added={len(skipped) - skipped_before_stage}",
|
||||
flush=True,
|
||||
)
|
||||
@@ -1391,7 +1458,7 @@ def run_submission(
|
||||
uniqueness_rejected_candidates: list[dict[str, Any]] = []
|
||||
target_submit_count = resolve_max_submit_count(
|
||||
args=args,
|
||||
planned_count=len(candidates),
|
||||
planned_count=candidate_model_count(candidates),
|
||||
remaining_daily_quota=remaining_daily_quota,
|
||||
)
|
||||
if strategy_manager is not None:
|
||||
@@ -1642,10 +1709,13 @@ def run_submission(
|
||||
"olderHistoryScanEnabled": allow_older_models_for_scan,
|
||||
},
|
||||
"scanLimit": scan_limit,
|
||||
"rawCandidateGoal": raw_candidate_goal,
|
||||
"candidatePoolLimit": candidate_pool_limit,
|
||||
"candidateGoal": candidate_goal,
|
||||
"scanStages": scan_stages,
|
||||
"scannedModels": scanned_model_count,
|
||||
"candidateCount": len(candidates),
|
||||
"candidateModelCount": candidate_model_count(candidates),
|
||||
"candidateFrameworkCounts": dict(candidate_framework_counts.most_common()),
|
||||
"targetSubmitCount": target_submit_count,
|
||||
"maxSubmitAttempts": 0 if args.dry_run else min(
|
||||
|
||||
Reference in New Issue
Block a user