964362818e
Do not cache transient repository-list failures as empty weight sets. Deduplicate normalized file paths without lowering reported file sizes. Record whitelisted OOM evidence and bounded discovery diagnostics.
511 lines
19 KiB
Python
511 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
from architecture_compatibility import architecture_compatibility_key, architecture_profiles
|
|
from llm_classifier import LLMAssistedClassifier
|
|
from models import ModelInspection
|
|
from common import parse_datetime, utc_now
|
|
|
|
|
|
# ModelHub allocations are preferred over product-card capacities. Nine values
|
|
# were repeated consistently across 299 structured PREFLIGHT_OOM reports pulled
|
|
# on 2026-08-10. The remaining currently verifiable devices use published card
|
|
# specifications and are replaced automatically when ModelHub returns its own
|
|
# observed allocation in a future failure log.
|
|
GPU_MEMORY_EVIDENCE: dict[str, dict[str, Any]] = {
|
|
"Ascend_910-b4": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:8"},
|
|
"Cambricon_mlu-370-x4": {"memoryGiB": 24.0, "source": "modelhub_preflight_oom:49"},
|
|
"Iluvatar_bi-100": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:48"},
|
|
"Iluvatar_bi-150": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:24"},
|
|
"Iluvatar_mrv-100": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:50"},
|
|
"MetaX_c-500": {"memoryGiB": 64.0, "source": "modelhub_preflight_oom:29"},
|
|
"Sunrise_pt-200-x1": {"memoryGiB": 64.0, "source": "modelhub_preflight_oom:36"},
|
|
"Vastai_va16": {"memoryGiB": 32.0, "source": "modelhub_preflight_oom:40"},
|
|
"hygon_k100-ai": {"memoryGiB": 64.0, "source": "modelhub_preflight_oom:15"},
|
|
"Ascend_910-b3": {
|
|
"memoryGiB": 64.0,
|
|
"source": "published_card_spec",
|
|
"sourceUrl": "https://aclanthology.org/2025.emnlp-main.1630.pdf",
|
|
},
|
|
"Biren_166m": {
|
|
"memoryGiB": 64.0,
|
|
"source": "manufacturer_spec",
|
|
"sourceUrl": "https://www.birentech.com/news/id6rz98v3obczy77cmxzfgk3/",
|
|
},
|
|
"Cambricon_mlu-370-x8": {
|
|
"memoryGiB": 48.0,
|
|
"source": "manufacturer_spec",
|
|
"sourceUrl": "https://cambricon.com/index.php?a=lists&c=index&catid=406&m=content",
|
|
},
|
|
"Kunlunxin_p-800": {
|
|
"memoryGiB": 96.0,
|
|
"source": "published_procurement_spec",
|
|
"sourceUrl": "https://pms2g.shrcb.com/cms/cmscaigougg/1d93b4b8166041e096c65d739073ade1.html",
|
|
},
|
|
"Mthreads_s4000": {
|
|
"memoryGiB": 48.0,
|
|
"source": "manufacturer_spec",
|
|
"sourceUrl": "https://docs.mthreads.com/s4000/s4000-doc-online/product_specifications/",
|
|
},
|
|
}
|
|
OBSERVED_GPU_MEMORY_GIB = {
|
|
gpu: float(evidence["memoryGiB"])
|
|
for gpu, evidence in GPU_MEMORY_EVIDENCE.items()
|
|
}
|
|
MODEL_LOAD_OVERHEAD = 1.20
|
|
TEXT_TOKENIZER_TASKS = {
|
|
"text-generation",
|
|
"visual-multi-modal",
|
|
"reinforcement_learning",
|
|
"question_answering",
|
|
"feature_emb",
|
|
"text_classification",
|
|
}
|
|
# A small, deliberately conservative baseline used only to decide whether an
|
|
# architecture needs semantic review. It is not an allowlist: unknown values
|
|
# remain eligible when the optional LLM is disabled or abstains.
|
|
MATURE_MODEL_TYPES = {
|
|
"baichuan",
|
|
"bert",
|
|
"bloom",
|
|
"chatglm",
|
|
"cohere",
|
|
"deepseek_v2",
|
|
"deepseek_v3",
|
|
"falcon",
|
|
"gemma",
|
|
"gemma2",
|
|
"gemma3",
|
|
"glm",
|
|
"glm4",
|
|
"gpt2",
|
|
"gpt_bigcode",
|
|
"internlm",
|
|
"internlm2",
|
|
"llama",
|
|
"mistral",
|
|
"mixtral",
|
|
"mpt",
|
|
"opt",
|
|
"phi",
|
|
"phi3",
|
|
"qwen2",
|
|
"qwen2_moe",
|
|
"qwen3",
|
|
"qwen3_moe",
|
|
"roberta",
|
|
"t5",
|
|
"whisper",
|
|
"xlm-roberta",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PreflightAssessment:
|
|
allowed: bool
|
|
config_params: str
|
|
reason: str | None
|
|
warnings: tuple[str, ...]
|
|
ambiguous_reasons: tuple[str, ...]
|
|
metadata: dict[str, Any]
|
|
|
|
|
|
class CandidatePreflightAdvisor:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
llm_classifier: LLMAssistedClassifier | None = None,
|
|
gpu_memory_gib: dict[str, float] | None = None,
|
|
) -> None:
|
|
self.llm_classifier = llm_classifier
|
|
environment_memory = _load_gpu_memory_overrides()
|
|
self.gpu_memory_gib = {
|
|
**OBSERVED_GPU_MEMORY_GIB,
|
|
**environment_memory,
|
|
**(gpu_memory_gib or {}),
|
|
}
|
|
self.gpu_memory_evidence = {
|
|
gpu: dict(GPU_MEMORY_EVIDENCE.get(gpu) or {"memoryGiB": memory, "source": "constructor_override"})
|
|
for gpu, memory in self.gpu_memory_gib.items()
|
|
}
|
|
for gpu in environment_memory:
|
|
self.gpu_memory_evidence[gpu] = {
|
|
"memoryGiB": self.gpu_memory_gib[gpu],
|
|
"source": "environment_override",
|
|
}
|
|
for gpu in (gpu_memory_gib or {}):
|
|
self.gpu_memory_evidence[gpu] = {
|
|
"memoryGiB": self.gpu_memory_gib[gpu],
|
|
"source": "constructor_override",
|
|
}
|
|
self._lock = threading.Lock()
|
|
self._assessed = 0
|
|
self._hard_blocks = 0
|
|
self._llm_blocks = 0
|
|
self._context_clamps = 0
|
|
self._ambiguous = 0
|
|
self._architecture_blocks_applied = 0
|
|
self._feedback_stats: dict[str, Any] = {}
|
|
self._architecture_compatibility_blocks: dict[str, dict[str, Any]] = {}
|
|
|
|
def set_feedback_stats(self, report: dict[str, Any] | None) -> None:
|
|
self._feedback_stats = report if isinstance(report, dict) else {}
|
|
raw_blocks = self._feedback_stats.get("architectureCompatibilityBlocks") or {}
|
|
self._architecture_compatibility_blocks = (
|
|
raw_blocks if isinstance(raw_blocks, dict) else {}
|
|
)
|
|
for gpu, value in (self._feedback_stats.get("observedGpuMemoryGiB") or {}).items():
|
|
try:
|
|
memory_gib = float(value)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if not 0 < memory_gib <= 1024:
|
|
continue
|
|
self.gpu_memory_gib[str(gpu)] = memory_gib
|
|
self.gpu_memory_evidence[str(gpu)] = {
|
|
"memoryGiB": memory_gib,
|
|
"source": "local_modelhub_preflight_oom",
|
|
}
|
|
|
|
def assess(
|
|
self,
|
|
*,
|
|
inspection: ModelInspection,
|
|
task_type: str,
|
|
target_gpu: str,
|
|
framework: str,
|
|
config_params: str,
|
|
) -> PreflightAssessment:
|
|
warnings: list[str] = []
|
|
ambiguous: list[str] = []
|
|
metadata: dict[str, Any] = {
|
|
"modelType": inspection.model_type,
|
|
"architectures": inspection.architectures,
|
|
"quantizationMethod": inspection.quantization_method,
|
|
"estimatedLoadBytes": inspection.estimated_load_bytes(framework),
|
|
"repositoryOnDiskBytes": inspection.repository_size_bytes,
|
|
"recursiveRepositoryBytes": inspection.exact_repository_size_bytes,
|
|
"publishedSizeBytes": inspection.published_size_bytes,
|
|
"gpuMemoryGiB": self.gpu_memory_gib.get(target_gpu),
|
|
"gpuMemoryEvidence": self.gpu_memory_evidence.get(target_gpu),
|
|
}
|
|
with self._lock:
|
|
self._assessed += 1
|
|
|
|
learned_block = self._matching_architecture_block(
|
|
inspection=inspection,
|
|
target_gpu=target_gpu,
|
|
framework=framework,
|
|
task_type=task_type,
|
|
)
|
|
if learned_block is not None:
|
|
metadata["architectureCompatibilityBlock"] = learned_block
|
|
with self._lock:
|
|
self._architecture_blocks_applied += 1
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_learned_architecture_incompatible",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
|
|
# Empty file_paths means an injected/test inspection lacks structural
|
|
# metadata. Real discoveries with an empty tree already fail the weight
|
|
# compatibility gate, so do not make this test/fallback state a blocker.
|
|
has_structure_metadata = bool(inspection.file_paths)
|
|
applies_text_structure_rules = (
|
|
task_type in TEXT_TOKENIZER_TASKS and framework != "llamacpp"
|
|
)
|
|
if has_structure_metadata and applies_text_structure_rules:
|
|
if not inspection.has_root_config:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_missing_root_config",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
if not inspection.has_root_standard_weights:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_missing_root_weights",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
if not inspection.has_root_tokenizer:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_missing_root_tokenizer",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
|
|
estimated_bytes = inspection.estimated_load_bytes(framework)
|
|
repository_bytes = inspection.repository_size_bytes
|
|
memory_gib = self.gpu_memory_gib.get(target_gpu)
|
|
if has_structure_metadata and memory_gib is None:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_gpu_memory_unknown",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
if memory_gib and has_structure_metadata and repository_bytes is None:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_model_size_unknown",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
memory_sizing_bytes = repository_bytes or estimated_bytes
|
|
if memory_sizing_bytes and memory_gib:
|
|
required_gib = memory_sizing_bytes / (1024**3) * MODEL_LOAD_OVERHEAD
|
|
metadata["estimatedRequiredGiB"] = round(required_gib, 3)
|
|
metadata["memorySizingBasis"] = (
|
|
"published_repository_upper_bound"
|
|
if (inspection.published_size_bytes or 0) > (inspection.exact_repository_size_bytes or 0)
|
|
else ("recursive_repository_on_disk" if repository_bytes else "selected_weights_fallback")
|
|
)
|
|
metadata["maximumRepositorySizeGiB"] = round(memory_gib / MODEL_LOAD_OVERHEAD, 3)
|
|
if required_gib > memory_gib:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_predicted_oom",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
|
|
profile_key = "|".join(
|
|
(
|
|
target_gpu,
|
|
framework,
|
|
task_type,
|
|
inspection.model_type or "unknown",
|
|
inspection.quantization_method or "none",
|
|
)
|
|
)
|
|
profile_feedback = (
|
|
(self._feedback_stats.get("recentProfileCombinationStats") or {}).get(profile_key)
|
|
or {}
|
|
)
|
|
if profile_feedback:
|
|
metadata["recentProfileFeedback"] = profile_feedback
|
|
last_terminal_at = parse_datetime(profile_feedback.get("lastTerminalAt"))
|
|
consecutive_failures = int(profile_feedback.get("consecutiveFailures") or 0)
|
|
circuit_open = bool(
|
|
last_terminal_at is not None
|
|
and consecutive_failures >= 5
|
|
and last_terminal_at + timedelta(hours=12) > utc_now()
|
|
)
|
|
if circuit_open:
|
|
return self._hard_block(
|
|
config_params,
|
|
"preflight_recent_profile_failure_circuit",
|
|
warnings,
|
|
ambiguous,
|
|
metadata,
|
|
)
|
|
|
|
rendered_config, context_clamped = clamp_context_length(
|
|
config_params,
|
|
inspection.max_context_length,
|
|
)
|
|
if context_clamped:
|
|
warnings.append("preflight_context_length_clamped")
|
|
with self._lock:
|
|
self._context_clamps += 1
|
|
|
|
if inspection.config_fetch_error:
|
|
warnings.append("model_config_metadata_unavailable")
|
|
elif applies_text_structure_rules:
|
|
if not inspection.model_type:
|
|
ambiguous.append("model_type_missing_or_unknown")
|
|
elif inspection.model_type.lower() not in MATURE_MODEL_TYPES:
|
|
ambiguous.append(f"architecture_not_in_mature_baseline:{inspection.model_type.lower()}")
|
|
auto_map = inspection.model_config.get("auto_map")
|
|
if isinstance(auto_map, dict) and auto_map:
|
|
ambiguous.append("custom_remote_code_architecture")
|
|
if inspection.quantization_method:
|
|
ambiguous.append(f"quantization_compatibility:{inspection.quantization_method}")
|
|
if inspection.architectures and any(
|
|
not architecture.endswith(
|
|
(
|
|
"ForCausalLM",
|
|
"ForConditionalGeneration",
|
|
"ForSequenceClassification",
|
|
"Model",
|
|
)
|
|
)
|
|
for architecture in inspection.architectures
|
|
):
|
|
ambiguous.append("nonstandard_architecture_name")
|
|
|
|
if ambiguous:
|
|
with self._lock:
|
|
self._ambiguous += 1
|
|
llm = self.llm_classifier
|
|
review_gate = getattr(llm, "should_review_candidate", None) if llm is not None else None
|
|
should_review = bool(
|
|
llm is not None
|
|
and llm.enabled
|
|
and (review_gate(ambiguous) if callable(review_gate) else True)
|
|
)
|
|
if should_review:
|
|
decision = llm.classify(
|
|
inspection=inspection,
|
|
task_type=task_type,
|
|
target_gpu=target_gpu,
|
|
framework=framework,
|
|
ambiguous_reasons=ambiguous,
|
|
)
|
|
metadata["llmDecision"] = decision
|
|
if llm.blocks(decision):
|
|
with self._lock:
|
|
self._llm_blocks += 1
|
|
reason = str(decision.get("reason") or "incompatible")
|
|
return PreflightAssessment(
|
|
allowed=False,
|
|
config_params=rendered_config,
|
|
reason=f"llm_high_confidence_incompatible:{reason}",
|
|
warnings=tuple(warnings),
|
|
ambiguous_reasons=tuple(ambiguous),
|
|
metadata=metadata,
|
|
)
|
|
warnings.append(f"llm_review_{decision.get('decision') or 'abstain'}")
|
|
elif llm is not None and llm.enabled:
|
|
warnings.append("llm_review_not_needed")
|
|
else:
|
|
warnings.append("llm_review_unavailable")
|
|
|
|
return PreflightAssessment(
|
|
allowed=True,
|
|
config_params=rendered_config,
|
|
reason=None,
|
|
warnings=tuple(warnings),
|
|
ambiguous_reasons=tuple(ambiguous),
|
|
metadata=metadata,
|
|
)
|
|
|
|
def _hard_block(
|
|
self,
|
|
config_params: str,
|
|
reason: str,
|
|
warnings: list[str],
|
|
ambiguous: list[str],
|
|
metadata: dict[str, Any],
|
|
) -> PreflightAssessment:
|
|
with self._lock:
|
|
self._hard_blocks += 1
|
|
return PreflightAssessment(
|
|
allowed=False,
|
|
config_params=config_params,
|
|
reason=reason,
|
|
warnings=tuple(warnings),
|
|
ambiguous_reasons=tuple(ambiguous),
|
|
metadata=metadata,
|
|
)
|
|
|
|
def _matching_architecture_block(
|
|
self,
|
|
*,
|
|
inspection: ModelInspection,
|
|
target_gpu: str,
|
|
framework: str,
|
|
task_type: str,
|
|
) -> dict[str, Any] | None:
|
|
for profile in architecture_profiles(inspection.model_type, inspection.architectures):
|
|
key = architecture_compatibility_key(
|
|
target_gpu,
|
|
framework,
|
|
task_type,
|
|
profile["signature"],
|
|
)
|
|
if key is None:
|
|
continue
|
|
block = self._architecture_compatibility_blocks.get(key)
|
|
if not isinstance(block, dict):
|
|
continue
|
|
expires_at = parse_datetime(block.get("expiresAt"))
|
|
if expires_at is None or expires_at <= utc_now():
|
|
continue
|
|
return dict(block)
|
|
return None
|
|
|
|
def summary(self) -> dict[str, Any]:
|
|
with self._lock:
|
|
summary = {
|
|
"enabled": True,
|
|
"assessedCandidates": self._assessed,
|
|
"hardBlocks": self._hard_blocks,
|
|
"llmBlocks": self._llm_blocks,
|
|
"ambiguousCandidates": self._ambiguous,
|
|
"architectureCompatibilityBlocksLoaded": len(
|
|
self._architecture_compatibility_blocks
|
|
),
|
|
"architectureCompatibilityBlocksApplied": self._architecture_blocks_applied,
|
|
"contextLengthClamps": self._context_clamps,
|
|
"knownGpuMemoryGiB": dict(self.gpu_memory_gib),
|
|
"gpuMemoryEvidence": dict(self.gpu_memory_evidence),
|
|
}
|
|
summary["llm"] = self.llm_classifier.summary() if self.llm_classifier else {"enabled": False}
|
|
return summary
|
|
|
|
|
|
def clamp_context_length(config: str, maximum: int | None) -> tuple[str, bool]:
|
|
if not maximum or maximum <= 0:
|
|
return config, False
|
|
changed = False
|
|
|
|
def replace_value(match: re.Match[str]) -> str:
|
|
nonlocal changed
|
|
current = int(match.group("value"))
|
|
if current <= maximum:
|
|
return match.group(0)
|
|
changed = True
|
|
return f"{match.group('prefix')}{maximum}{match.group('suffix')}"
|
|
|
|
patterns = (
|
|
r"(?P<prefix>(?:max_model_len|max_seq_len)\s*[:=]\s*['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
|
|
r"(?P<prefix>['\"]max_seq_len['\"]\s*:\s*['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
|
|
r"(?P<prefix>MAX_MODEL_LEN\s*,?\s*value\s*:\s*['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
|
|
r"(?P<prefix>--max-model-len(?:\s+|\s*,\s*(?:\n\s*)?|\s*\n\s*-\s*)['\"]?)(?P<value>\d+)(?P<suffix>['\"]?)",
|
|
)
|
|
rendered = config
|
|
for pattern in patterns:
|
|
rendered = re.sub(pattern, replace_value, rendered, flags=re.IGNORECASE)
|
|
return rendered, changed
|
|
|
|
|
|
def _load_gpu_memory_overrides() -> dict[str, float]:
|
|
raw = os.getenv("MODELHUB_GPU_MEMORY_GIB_JSON")
|
|
if not raw:
|
|
return {}
|
|
try:
|
|
payload = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
if not isinstance(payload, dict):
|
|
return {}
|
|
result: dict[str, float] = {}
|
|
for gpu, value in payload.items():
|
|
try:
|
|
parsed = float(value)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if parsed > 0:
|
|
result[str(gpu)] = parsed
|
|
return result
|