126 lines
5.1 KiB
Python
126 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from common import read_jsonl
|
|
from gpu_strategy import wilson_lower_bound
|
|
|
|
|
|
class SafeConfigOptimizer:
|
|
"""Apply only repeatedly successful, sanitized parameter vectors."""
|
|
|
|
def __init__(self, *, intents_path: Path | str, outcomes_path: Path | str) -> None:
|
|
self.intents_path = Path(intents_path)
|
|
self.outcomes_path = Path(outcomes_path)
|
|
self._qualified = self._learn()
|
|
|
|
def _learn(self) -> dict[tuple[str, str, str], list[dict[str, Any]]]:
|
|
recent_path = Path(".modelhub_state/recent_outcomes.jsonl")
|
|
outcomes_by_key: dict[str, dict[str, Any]] = {}
|
|
for item in [*read_jsonl(recent_path), *read_jsonl(self.outcomes_path)]:
|
|
task_id = str(item.get("taskId") or "")
|
|
if task_id:
|
|
outcomes_by_key[task_id] = item
|
|
outcomes = list(outcomes_by_key.values())
|
|
by_task = {
|
|
str(item.get("taskId")): item
|
|
for item in outcomes
|
|
if item.get("taskId") is not None and item.get("outcome") in {"success", "failed"}
|
|
}
|
|
groups: dict[tuple[str, str, str, str], dict[str, Any]] = defaultdict(
|
|
lambda: {"success": 0, "failure": 0, "models": set(), "vector": {}}
|
|
)
|
|
for intent in read_jsonl(self.intents_path):
|
|
task_id = str(intent.get("taskId") or "")
|
|
outcome = by_task.get(task_id)
|
|
vector = intent.get("safeConfigVector")
|
|
if outcome is None or not isinstance(vector, dict) or not vector:
|
|
continue
|
|
route = (
|
|
str(intent.get("taskType") or ""),
|
|
str(intent.get("targetGpu") or ""),
|
|
str(intent.get("framework") or ""),
|
|
str(intent.get("configFingerprint") or ""),
|
|
)
|
|
group = groups[route]
|
|
group["vector"] = dict(vector)
|
|
group["models"].add(str(intent.get("repoId") or ""))
|
|
if outcome.get("outcome") == "success":
|
|
group["success"] += 1
|
|
elif not outcome.get("platformFailure") and not outcome.get("policyCancelled"):
|
|
group["failure"] += 1
|
|
|
|
qualified: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list)
|
|
for (task_type, gpu, framework, fingerprint), group in groups.items():
|
|
success = int(group["success"])
|
|
failure = int(group["failure"])
|
|
if success < 5 or failure > 0 or len(group["models"]) < 2:
|
|
continue
|
|
qualified[(task_type, gpu, framework)].append(
|
|
{
|
|
"fingerprint": fingerprint,
|
|
"success": success,
|
|
"total": success + failure,
|
|
"lowerBound": wilson_lower_bound(success, success + failure),
|
|
"vector": group["vector"],
|
|
}
|
|
)
|
|
return dict(qualified)
|
|
|
|
def optimize(
|
|
self,
|
|
*,
|
|
task_type: str,
|
|
target_gpu: str,
|
|
framework: str,
|
|
official_config: str,
|
|
official_lower_bound: float,
|
|
) -> tuple[str, dict[str, Any]]:
|
|
choices = self._qualified.get((task_type, target_gpu, framework)) or []
|
|
choices = [
|
|
item
|
|
for item in choices
|
|
if float(item.get("lowerBound") or 0.0) >= float(official_lower_bound) + 0.05
|
|
]
|
|
if not choices:
|
|
return official_config, {"source": "official", "applied": False}
|
|
choice = max(choices, key=lambda item: (float(item["lowerBound"]), int(item["success"])))
|
|
patched = official_config
|
|
vector = choice["vector"]
|
|
substitutions = {
|
|
"gpuNum": (r"(\bgpu_num\s*:\s*['\"]?)\d+", r"\g<1>{}"),
|
|
"tensorParallel": (
|
|
r"((?:--tensor-parallel-size|-tp)\s*[, ]?\s*['\"]?)\d+",
|
|
r"\g<1>{}",
|
|
),
|
|
"maxModelLen": (
|
|
r"((?:--max-model-len|max_model_len|max_seq_len)\s*[: ,]+\s*['\"]?)\d+",
|
|
r"\g<1>{}",
|
|
),
|
|
"gpuMemoryUtilization": (
|
|
r"((?:--gpu-memory-utilization|gpu_memory_utilization)\s*[: ,]+\s*['\"]?)[0-9.]+",
|
|
r"\g<1>{}",
|
|
),
|
|
}
|
|
for key, (pattern, replacement) in substitutions.items():
|
|
if key in vector:
|
|
patched = re.sub(pattern, replacement.format(vector[key]), patched, flags=re.IGNORECASE)
|
|
for key, option in (("dtype", "dtype"), ("quantization", "quantization"), ("loadFormat", "load[_-]format")):
|
|
if key in vector:
|
|
patched = re.sub(
|
|
rf"((?:--{option}|{option})\s*[: ,]+\s*['\"]?)[A-Za-z0-9_-]+",
|
|
rf"\g<1>{vector[key]}",
|
|
patched,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
return patched, {
|
|
"source": "official_plus_learned_patch",
|
|
"applied": patched != official_config,
|
|
"evidenceSuccess": choice["success"],
|
|
"evidenceLowerBound": choice["lowerBound"],
|
|
"fingerprint": choice["fingerprint"],
|
|
}
|