507 lines
18 KiB
Python
507 lines
18 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,
|
|
"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"] = (
|
|
"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
|