hengyan/vllm_agent_strategy
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Initial vLLM agent strategy

Browse Source
This commit is contained in:
hengyan
2026-07-21 16:04:18 +08:00
commit d17bb21bbb
48 changed files with 3006 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
app/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""vLLM intelligent-agent strategy package."""

1
app/clients/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""API clients."""

153
app/clients/modelhub.py Normal file
View File

@@ -0,0 +1,153 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import logging
import requests
from app.settings import Settings
LOG = logging.getLogger(__name__)
@dataclass
class ApiResult:
result: str
code: str | None = None
message: str | None = None
payload: dict[str, Any] | None = None
class ModelHubClient:
def __init__(self, settings: Settings):
self.settings = settings
self.session = requests.Session()
def _headers(self) -> dict[str, str]:
return {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.settings.auth_token}",
"Xc-Token": self.settings.auth_token,
}
def is_model_absent_from_modelhub(self, model_id: str) -> bool:
url = self.settings.modelhub_api_base + "/computility/models/list/page/vo"
data = {"current": 1, "pageSize": 20, "searchText": model_id}
try:
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
response.raise_for_status()
total = int(response.json().get("data", {}).get("total", 0))
return total == 0
except Exception as exc:
LOG.warning("modelhub existence check failed for %s: %s", model_id, exc)
return True
def list_modelhub_models(self, current: int = 1, page_size: int = 100, search_text: str = "") -> dict[str, Any]:
url = self.settings.modelhub_api_base + "/computility/models/list/page/vo"
data = {"current": current, "pageSize": page_size, "searchText": search_text}
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
response.raise_for_status()
body = response.json()
if body.get("code") != 0:
raise RuntimeError(body.get("message") or body.get("msg") or "list_modelhub_models failed")
return body.get("data") or {}
def list_success_download_tasks(self, current: int = 1, page_size: int = 50) -> dict[str, Any]:
url = self.settings.modelhub_adminapi_base + "/async/task/model-download-task"
params = {"current": current, "pageSize": page_size, "status": "SUCCESS"}
response = self.session.get(url, headers=self._headers(), params=params, timeout=self.settings.request_timeout_seconds)
response.raise_for_status()
body = response.json()
if body.get("code") != 0:
raise RuntimeError(body.get("message") or "list_success_download_tasks failed")
return body.get("data") or {}
def sync_from_huggingface(self, model_id: str) -> ApiResult:
url = self.settings.modelhub_adminapi_base + "/computility/models/sync-from-hugging-face"
data = {"modelId": model_id, "forceUpdate": False, "operatorEmail": self.settings.email}
try:
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
payload = response.json()
return ApiResult("success" if payload.get("code") == 0 else "failed", str(payload.get("code")), payload.get("message"), payload)
except Exception as exc:
return ApiResult("failed", None, f"sync exception: {exc}", None)
def create_download_task(self, model_id: str, source: str = "HUGGING_FACE", weight_file: str | None = None) -> ApiResult:
url = self.settings.modelhub_url + "/adminApi/async/task/model-download-task"
data: dict[str, Any] = {
"hfToken": self.settings.hf_token,
"modelId": model_id,
"source": source,
"stillDownloadAlreadySuccessDownloadedModel": False,
}
if weight_file:
data["allowPatterns"] = [weight_file]
try:
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
payload = response.json()
code = payload.get("code")
msg = payload.get("message", "")
if code == 0:
result = "created"
elif code == 40000:
result = "running_elsewhere"
elif code == 60004:
result = "already_downloaded"
elif code == 60005:
result = "download_pool_full"
else:
result = "failed"
return ApiResult(result, str(code), msg, payload)
except Exception as exc:
return ApiResult("failed", None, f"download exception: {exc}", None)
def get_download_status(self, model_id: str) -> str:
url = self.settings.modelhub_url + "/adminApi/async/task/model-download-task"
params = {"modelId": model_id, "userEmail": self.settings.email, "userId": self.settings.user_id}
try:
response = self.session.get(url, headers=self._headers(), params=params, timeout=self.settings.request_timeout_seconds)
response.raise_for_status()
records = response.json().get("data", {}).get("records", [])
if records:
return records[0].get("status") or "None"
return "None"
except Exception as exc:
LOG.warning("download status check failed for %s: %s", model_id, exc)
return "QUERY_FAILED"
def build_contest_payload(self, model_id: str, gpu_type: str, config: str) -> dict[str, Any]:
payload: dict[str, Any] = {
"contestApiToken": self.settings.contest_api_token,
"contributors": self.settings.contributors,
"gpuTypes": [gpu_type],
"modelId": model_id,
"submissionConfig": [{"gpuType": gpu_type, "config": config}],
"taskType": "text-generation",
}
if self.settings.inject_strategy_id:
payload[self.settings.contest_task_strategy_field] = self.settings.strategy_id
return payload
def create_contest_task(self, model_id: str, gpu_type: str, config: str) -> ApiResult:
payload = self.build_contest_payload(model_id, gpu_type, config)
if self.settings.submit_dry_run:
LOG.info("dry-run submit payload: model=%s gpu=%s strategy_field=%s", model_id, gpu_type, self.settings.contest_task_strategy_field)
return ApiResult("success", "DRY_RUN", "dry run", {"payload": payload})
url = self.settings.modelhub_adminapi_base + "/async/task/create-contest-task"
try:
response = self.session.post(url, headers=self._headers(), json=payload, timeout=self.settings.request_timeout_seconds)
body = response.json()
code = body.get("code")
message = body.get("message", "")
if code == 0 and message == "ok":
result = "success"
elif code == 40000 and "正在验证中" in message:
result = "conflict"
elif code == 60007:
result = "queue_full"
else:
result = "failed"
return ApiResult(result, str(code), message, body)
except Exception as exc:
return ApiResult("failed", None, f"submit exception: {exc}", None)

1
app/crawlers/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Candidate crawlers."""

85
app/crawlers/download_success.py Normal file
View File

@@ -0,0 +1,85 @@
from __future__ import annotations
import logging
from app.clients.modelhub import ModelHubClient
from app.domain.model_type import engine_for_model
from app.domain.priorities import PRIORITY_OTHERS_DOWNLOADED
from app.settings import Settings
from app.storage.repositories import Repository
LOG = logging.getLogger(__name__)
SOURCE_MAP = {
"HuggingFace": "HUGGING_FACE",
"HUGGING_FACE": "HUGGING_FACE",
"ModelScope": "MODEL_SCOPE",
"MODEL_SCOPE": "MODEL_SCOPE",
}
def normalize_source(raw: str | None) -> str:
if not raw:
return "HUGGING_FACE"
return SOURCE_MAP.get(raw, raw)
class DownloadSuccessCrawler:
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
self.settings = settings
self.repo = repo
self.client = client
def refresh(self) -> int:
if not self.settings.enable_download_success_crawler:
return 0
imported = 0
seen: set[str] = set()
page = 1
while True:
data = self.client.list_success_download_tasks(page, self.settings.download_success_page_size)
records = data.get("records") or []
total = int(data.get("total") or 0)
if not records:
break
for record in records:
args = record.get("args") or {}
model_id = args.get("model_id") or args.get("modelId")
if not model_id or model_id in seen:
continue
seen.add(model_id)
source = normalize_source(args.get("source"))
engine = engine_for_model(model_id)
if engine != "vllm":
self.repo.upsert_candidate(
model_id=model_id,
source=source,
origin="model_download_success_page_gguf",
priority=PRIORITY_OTHERS_DOWNLOADED,
known_downloaded_by_others=True,
self_download_allowed=False,
notes=f"download task {record.get('id')} SUCCESS; GGUF/llamacpp candidate",
)
else:
self.repo.upsert_candidate(
model_id=model_id,
source=source,
origin="model_download_success_page",
priority=PRIORITY_OTHERS_DOWNLOADED,
known_downloaded_by_others=True,
self_download_allowed=False,
notes=f"download task {record.get('id')} SUCCESS",
)
imported += 1
if self.settings.download_success_max_pages and page >= self.settings.download_success_max_pages:
break
if total and page * self.settings.download_success_page_size >= total:
break
if len(records) < self.settings.download_success_page_size:
break
page += 1
LOG.info("download-success crawler imported/updated %s candidates", imported)
return imported

109
app/crawlers/not_adapted.py Normal file
View File

@@ -0,0 +1,109 @@
from __future__ import annotations
import logging
from app.clients.modelhub import ModelHubClient
from app.domain.machines import machine_to_gpu_alias, parse_machine_names
from app.domain.model_type import engine_for_model
from app.domain.priorities import PRIORITY_BOUNTY
from app.settings import Settings
from app.storage.repositories import Repository
LOG = logging.getLogger(__name__)
def _model_task_name(model: dict) -> str:
task_info = model.get("taskLevelsInfo") or {}
return model.get("taskLevelChineseName") or task_info.get("taskLevelChineseName") or ""
def _adapted_machine_names(model: dict) -> set[str]:
machines = model.get("machines") or []
return {m.get("machineName") for m in machines if m.get("machineName")}
class NotAdaptedCrawler:
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
self.settings = settings
self.repo = repo
self.client = client
def refresh(self) -> int:
if not self.settings.enable_not_adapted_crawler:
return 0
machine_names = parse_machine_names(self.settings.target_machine_names)
if not machine_names:
LOG.info("not-adapted crawler skipped: TARGET_MACHINE_NAMES is empty")
return 0
machine_alias_pairs: list[tuple[str, str]] = []
for machine_name in machine_names:
alias = machine_to_gpu_alias(machine_name)
if not alias:
LOG.warning("unknown target machine name, skipped: %s", machine_name)
continue
machine_alias_pairs.append((machine_name, alias))
if not machine_alias_pairs:
return 0
imported = 0
page = 1
while True:
data = self.client.list_modelhub_models(page, self.settings.modelhub_page_size, "")
records = data.get("records") or []
total = int(data.get("total") or 0)
if not records:
break
for model in records:
model_id = model.get("modelId")
if not model_id:
continue
task_name = _model_task_name(model)
if self.settings.target_task_level and self.settings.target_task_level not in task_name:
continue
adapted = _adapted_machine_names(model)
engine = engine_for_model(model_id)
for machine_name, alias in machine_alias_pairs:
if machine_name in adapted:
continue
if engine != "vllm":
self.repo.upsert_candidate(
model_id=model_id,
source="HUGGING_FACE",
origin="bounty_target_not_adapted_gguf",
priority=PRIORITY_BOUNTY,
is_bounty=True,
is_promote=True,
known_downloaded_by_others=True,
self_download_allowed=False,
target_gpu_aliases=alias,
notes=f"GGUF/llamacpp candidate for {machine_name}; vLLM phase skips submit",
)
imported += 1
continue
self.repo.upsert_candidate(
model_id=model_id,
source="HUGGING_FACE",
origin="bounty_target_not_adapted",
priority=PRIORITY_BOUNTY,
is_bounty=True,
is_promote=True,
known_downloaded_by_others=True,
self_download_allowed=False,
target_gpu_aliases=alias,
notes=f"not adapted on {machine_name}",
)
imported += 1
if self.settings.crawler_max_pages and page >= self.settings.crawler_max_pages:
break
if total and page * self.settings.modelhub_page_size >= total:
break
if len(records) < self.settings.modelhub_page_size:
break
page += 1
LOG.info("not-adapted crawler imported/updated %s candidates", imported)
return imported

1
app/domain/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Domain helpers for model submission strategy."""

41
app/domain/gpu.py Normal file
View File

@@ -0,0 +1,41 @@
from __future__ import annotations
VLLM_GPU_ALIASES = {
"910b": ["Ascend_910-b3", "Ascend_910-b4"],
"s4000": ["Mthreads_s4000"],
"k100": ["hygon_k100-ai"],
"mrv100": ["Iluvatar_mrv-100"],
"bi100": ["Iluvatar_bi-100"],
"bi150": ["Iluvatar_bi-150"],
"c500": ["MetaX_c-500"],
"mlu370-x4": ["Cambricon_mlu-370-x4"],
"mlu370-x8": ["Cambricon_mlu-370-x8"],
"p800": ["Kunlunxin_p-800"],
"166m": ["Biren_166m"],
"biren166m": ["Biren_166m"], # backward-compatible alias from the old config.py
}
# Sunrise s2 currently uses SGLang in the old config, so vLLM-only v1 skips it.
UNSUPPORTED_VLLM_ALIASES = {"s2"}
def parse_aliases(raw: str | None, defaults: list[str]) -> list[str]:
if not raw:
return defaults
return [item.strip() for item in raw.split(",") if item.strip()]
def expand_gpu_alias(alias: str) -> list[str]:
return VLLM_GPU_ALIASES.get(alias, [])
def expand_gpu_aliases(aliases: list[str]) -> list[tuple[str, str]]:
pairs: list[tuple[str, str]] = []
for alias in aliases:
for gpu_type in expand_gpu_alias(alias):
pairs.append((alias, gpu_type))
return pairs
def is_supported_alias(alias: str) -> bool:
return alias in VLLM_GPU_ALIASES

32
app/domain/machines.py Normal file
View File

@@ -0,0 +1,32 @@
from __future__ import annotations
MACHINE_TO_GPU_ALIAS = {
"MTT S4000": "s4000",
"Mthreads S4000": "s4000",
"摩尔线程 S4000": "s4000",
"Hygon K100": "k100",
"海光 K100": "k100",
"Kunlunxin P800": "p800",
"Kunlunxin P-800": "p800",
"昆仑芯 P800": "p800",
"Biren 166M": "166m",
"Biren 166m": "166m",
"壁砺 166M": "166m",
"壁砺166M": "166m",
"Ascend 910B": "910b",
"昇腾 910B": "910b",
"Iluvatar BI100": "bi100",
"Iluvatar BI150": "bi150",
"MetaX C500": "c500",
"Iluvatar MRV100": "mrv100",
"Cambricon MLU370-X4": "mlu370-x4",
"Cambricon MLU370-X8": "mlu370-x8",
}
def machine_to_gpu_alias(machine_name: str) -> str | None:
return MACHINE_TO_GPU_ALIAS.get(machine_name.strip())
def parse_machine_names(raw: str) -> list[str]:
return [item.strip() for item in raw.split(",") if item.strip()]

12
app/domain/model_type.py Normal file
View File

@@ -0,0 +1,12 @@
from __future__ import annotations
def is_gguf_model(model_id: str, extra_text: str = "") -> bool:
haystack = f"{model_id} {extra_text}".lower()
return "gguf" in haystack or ".gguf" in haystack
def engine_for_model(model_id: str, extra_text: str = "") -> str:
if is_gguf_model(model_id, extra_text):
return "llamacpp_candidate"
return "vllm"

16
app/domain/priorities.py Normal file
View File

@@ -0,0 +1,16 @@
PRIORITY_BOUNTY = 0
PRIORITY_PROMOTE = 10
PRIORITY_OTHERS_DOWNLOADED = 20
PRIORITY_SELF_DOWNLOAD = 30
PRIORITY_LOW_CONFIDENCE = 90
ORIGIN_TO_PRIORITY = {
"manual_bounty": PRIORITY_BOUNTY,
"bounty": PRIORITY_BOUNTY,
"manual_promote": PRIORITY_PROMOTE,
"promote": PRIORITY_PROMOTE,
"manual_downloaded_by_others": PRIORITY_OTHERS_DOWNLOADED,
"downloaded_by_others": PRIORITY_OTHERS_DOWNLOADED,
"hf_seed": PRIORITY_SELF_DOWNLOAD,
"self_download": PRIORITY_SELF_DOWNLOAD,
}

449
app/domain/vllm_configs.py Normal file
View File

@@ -0,0 +1,449 @@
from __future__ import annotations
# Templates are copied from config部分更新.py and use __MODEL_ID__ tokens instead of
# str.format placeholders so YAML inline maps like {name: ..., value: ...} stay intact.
VLLM_MRV100_CONFIG = """docker_image: harbor.4pd.io/hardcore-tech/iluvatar/llm-infer-iluvatar-mr:v0
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
framework: vllm
storage: gpfs
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
max_model_len: 4096
lang: en
temperature: 0.4
repetition_penalty: 1.1
top_p: 0.9
sut_config:
gpu_num: 1
values:
command: [vllm, serve, /model, --port, '20644', --served-model-name, llm, --max-model-len,
'4096', --gpu-memory-utilization, '0.9', --enforce-eager, --trust-remote-code,
-tp, '1']
ref_config:
gpu_num: 1
values:
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
'4096', --enforce-eager, --trust-remote-code, -tp, '1']
"""
VLLM_S4000_CONFIG = """docker_image: git.modelhub.org.cn:9443/enginex-mthreads/vllm-musa-qy2-py310:v0.8.4-release
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
framework: vllm
storage: gpfs
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
api: completion
lang: en
max_model_len: 4096
max_tokens: 1024
temperature: 0.7
repetition_penalty: 1.1
top_p: 0.9
sut_config:
gpu_num: 1
values:
command:
- vllm
- serve
- /model
- --port
- '8000'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --dtype
- auto
- --gpu-memory-utilization
- '0.95'
- -tp
- '1'
- --enforce-eager
- --trust-remote-code
ref_config:
gpu_num: 1
values:
command:
- vllm
- serve
- /model
- --port
- '8000'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --dtype
- auto
- --gpu-memory-utilization
- '0.95'
- -tp
- '1'
- --enforce-eager
- --trust-remote-code
"""
VLLM_910B_CONFIG = """
framework: vllm
docker_image: git.modelhub.org.cn:9443/enginex-ascend/vllm-ascend:v0.11.0rc0
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
storage: gpfs
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
api: completion
max_tokens: 1024
temperature: 0.7
repetition_penalty: 1.2
top_p: 0.9
lang: zh
max_model_len: 2048
sut_config:
gpu_num: 1
values:
command:
- vllm
- serve
- /model
- --port
- '8000'
- --served-model-name
- llm
- --max-model-len
- '2048'
- --dtype
- auto
- --gpu-memory-utilization
- '0.95'
- -tp
- '1'
- --enforce-eager
- --trust-remote-code
ref_config:
gpu_num: 1
values:
command:
- vllm
- serve
- /model
- --port
- '8000'
- --served-model-name
- llm
- --max-model-len
- '2048'
- --dtype
- auto
- --gpu-memory-utilization
- '0.95'
- -tp
- '1'
- --enforce-eager
- --trust-remote-code
"""
VLLM_BI100_CONFIG = """docker_image: git.modelhub.org.cn:9443/enginex-iluvatar-bi100/vllm:0.6.3
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
framework: vllm
storage: gpfs
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
lang: en
max_model_len: 4096
sut_config:
gpu_num: 1
environment:
MAX_MODEL_LEN: "4096"
values:
command: [/workspace/launch_service, --port, '80']
ref_config:
gpu_num: 1
environment:
MAX_MODEL_LEN: "4096"
values:
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
'4096', --enforce-eager, --trust-remote-code, -tp, '1']
"""
VLLM_BI150_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/sunruoxi/enginex-iluvatar-bi150-vllm-fix-tokenizer:v0.8.3
nv_docker_image: harbor-contest.4pd.io/sunruoxi/vllm-openai-fix-tokenizer:v0.11.0
framework: vllm_fix_tokenizer
nv_framework: vllm_fix_tokenizer
lang: en
max_model_len: 4096
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
sut_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- /model
- --port
- '80'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
ref_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- /model
- --port
- '80'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
"""
VLLM_C500_CONFIG = """docker_image: git.modelhub.org.cn:9443/enginex-metax/vllm:0.9.1
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
framework: vllm
storage: gpfs
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
api: completion
lang: en
max_model_len: 4096
max_tokens: 1024
temperature: 0.7
repetition_penalty: 1.1
top_p: 0.9
sut_config:
gpu_num: 1
values:
command: [/opt/conda/bin/vllm, serve, /model, --port, '20644', --served-model-name,
llm, --max-model-len, '4096', --gpu-memory-utilization, '0.9', -tp, '1', --enforce-eager,
--trust-remote-code]
ref_config:
gpu_num: 1
values:
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
'4096', -tp, '1', --enforce-eager, --trust-remote-code]
"""
VLLM_K100_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/i-peixingyu/k100-vllm-patched-v2.0:v2.0.0
nv_docker_image: harbor-contest.4pd.io/sunruoxi/vllm-openai-fix-tokenizer:v0.11.0
framework: vllm-patch-tokenizer
nv_framework: vllm_fix_tokenizer
max_model_len: 4096
api: completion
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
sut_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- --port
- '20644'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
ref_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- /model
- --port
- '80'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
"""
VLLM_MLU370_X4_CONFIG = """docker_image: harbor.4pd.io/hardcore-tech/cambricon-mlu370-pytorch:v25.01-torch2.5.0-torchmlu1.24.1-ubuntu22.04-py310
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
framework: vllm-mlu
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
sut_config:
values:
gpu_num: 1
env:
- {name: MAX_MODEL_LEN, value: 4096}
command: [vllm, serve, /model, --port, '8000', --served-model-name, llm, --max-model-len,
'4096', --trust-remote-code, --dtype, float16]
ref_config:
values:
cpu_num: 2
gpu_num: 1
env:
- {name: MAX_MODEL_LEN, value: 4096}
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
'4096', --trust-remote-code, --dtype, float16]
"""
VLLM_MLU370_X8_CONFIG = """docker_image: harbor.4pd.io/hardcore-tech/cambricon-mlu370-pytorch:v25.01-20260204
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
framework: vllm-mlu
api: completion
lang: en
temperature: 0.4
repetition_penalty: 1.1
top_p: 0.9
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
sut_config:
gpu_num: 1
values:
env:
- {name: MAX_MODEL_LEN, value: 4096}
- {name: VLLM_ALLOW_LONG_MAX_MODEL_LEN, value: '1'}
command: [vllm, serve, /model, --port, '8000', --served-model-name, llm, --max-model-len,
'4096', --trust-remote-code, --dtype, float16]
ref_config:
cpu_num: 2
gpu_num: 1
values:
env:
- {name: MAX_MODEL_LEN, value: 4096}
- {name: VLLM_ALLOW_LONG_MAX_MODEL_LEN, value: '1'}
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
'4096', --trust-remote-code, --dtype, float16]
"""
VLLM_P800_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/zheng/vllm-kunlunxin-p-800-tokenizer-patch:v1.0
nv_docker_image: harbor-contest.4pd.io/sunruoxi/vllm-openai-fix-tokenizer:v0.11.0
framework: vllm_tokenizer_patch
nv_framework: vllm_fix_tokenizer
api: completion
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
max_model_len: 4096
sut_config:
gpu_num: 1
values:
command:
- vllm
- serve
- /model
- --port
- '8000'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --gpu-memory-utilization
- '0.9'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
ref_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- /model
- --port
- '80'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
"""
VLLM_166M_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/sunruoxi/enginex-llm-biren166m-fix-tokenizer:v26.01
nv_docker_image: harbor.4pd.io/modelhubxc/enginex-nvidia/vllm:0.11.0-patch-tokenizer
modelhub_options:
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
mountPoint: /model
framework: vllm_fix_tokenizer
nv_framework: vllm_fix_tokenizer
lang: en
max_model_len: 4096
sut_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- /model
- --port
- '80'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
ref_config:
gpu_num: 1
values:
command:
- /opt/entrypoint.sh
- /model
- --port
- '80'
- --served-model-name
- llm
- --max-model-len
- '4096'
- --enforce-eager
- --trust-remote-code
- -tp
- '1'
"""
VLLM_GPU_CONFIG_DICT = {
"Ascend_910-b3": VLLM_910B_CONFIG,
"Ascend_910-b4": VLLM_910B_CONFIG,
"Mthreads_s4000": VLLM_S4000_CONFIG,
"hygon_k100-ai": VLLM_K100_CONFIG,
"Iluvatar_mrv-100": VLLM_MRV100_CONFIG,
"Iluvatar_bi-100": VLLM_BI100_CONFIG,
"Iluvatar_bi-150": VLLM_BI150_CONFIG,
"MetaX_c-500": VLLM_C500_CONFIG,
"Cambricon_mlu-370-x4": VLLM_MLU370_X4_CONFIG,
"Cambricon_mlu-370-x8": VLLM_MLU370_X8_CONFIG,
"Kunlunxin_p-800": VLLM_P800_CONFIG,
"Biren_166m": VLLM_166M_CONFIG,
}
def supported_gpu_types() -> set[str]:
return set(VLLM_GPU_CONFIG_DICT)
def gen_vllm_config(gpu_type: str, model_id: str, max_model_len: int = 1024) -> str:
if gpu_type not in VLLM_GPU_CONFIG_DICT:
raise ValueError(f"Unsupported vLLM gpu_type: {gpu_type}")
return VLLM_GPU_CONFIG_DICT[gpu_type].replace("__MODEL_ID__", model_id)

58
app/main.py Normal file
View File

@@ -0,0 +1,58 @@
from __future__ import annotations
import logging
import threading
from flask import Flask, jsonify
from app.clients.modelhub import ModelHubClient
from app.scheduler.loop import StrategyLoop
from app.settings import load_settings
from app.shutdown import install_signal_handlers, stop_event
from app.storage.db import connect
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
LOG = logging.getLogger(__name__)
def create_app() -> Flask:
app = Flask(__name__)
@app.get("/health")
def health():
return jsonify({"status": "ok"}), 200
@app.get("/status")
def status():
return jsonify({"status": "ok", "stopping": stop_event.is_set()}), 200
return app
def main():
install_signal_handlers()
settings = load_settings(require_secrets=True)
LOG.info(
"starting vLLM agent strategy dry_run=%s inject_strategy_id=%s strategy_field=%s db=%s",
settings.submit_dry_run,
settings.inject_strategy_id,
settings.contest_task_strategy_field,
settings.db_path,
)
conn = connect(settings.db_path)
client = ModelHubClient(settings)
loop = StrategyLoop(settings, __import__("app.storage.repositories", fromlist=["Repository"]).Repository(conn), client)
worker = threading.Thread(target=loop.run_forever, args=(stop_event,), name="strategy-loop", daemon=True)
worker.start()
app = create_app()
try:
app.run(host=settings.host, port=settings.port, threaded=True)
finally:
stop_event.set()
worker.join(timeout=25)
conn.close()
LOG.info("shutdown complete")
if __name__ == "__main__":
main()

1
app/scheduler/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Scheduler components."""

13
app/scheduler/budget.py Normal file
View File

@@ -0,0 +1,13 @@
from app.storage.repositories import Repository
class Budget:
def __init__(self, repo: Repository, max_tasks: int):
self.repo = repo
self.max_tasks = max_tasks
def remaining(self) -> int:
return max(0, self.max_tasks - self.repo.count_budgeted_tasks())
def can_create_task(self) -> bool:
return self.remaining() > 0

53
app/scheduler/downloader.py Normal file
View File

@@ -0,0 +1,53 @@
from app.clients.modelhub import ModelHubClient
from app.storage.repositories import Repository, future_seconds
class Downloader:
def __init__(self, repo: Repository, client: ModelHubClient, max_self_downloads: int, poll_interval_seconds: int):
self.repo = repo
self.client = client
self.max_self_downloads = max_self_downloads
self.poll_interval_seconds = poll_interval_seconds
def poll_active_downloads(self) -> int:
changed = 0
for download in self.repo.active_downloads_due(limit=20):
status = self.client.get_download_status(download["model_id"])
active = status in {"WAITING", "RUNNING", "QUERY_FAILED"}
self.repo.upsert_download(
download["model_id"],
download["source"],
status,
"self",
active_self_slot=active,
next_poll_at=future_seconds(self.poll_interval_seconds) if active else None,
)
changed += 1
return changed
def maybe_start_self_downloads(self, limit: int = 5) -> int:
started = 0
for task in self.repo.due_tasks(("waiting_download",), limit=limit):
if self.repo.count_active_self_downloads() >= self.max_self_downloads:
self.repo.mark_task_status(task["id"], "waiting_download", next_attempt_at=future_seconds(300))
continue
sync_result = self.client.sync_from_huggingface(task["model_id"])
if sync_result.result != "success":
self.repo.mark_task_status(task["id"], "waiting_download", sync_result.code, sync_result.message, future_seconds(600), increment_attempt=True)
continue
dl = self.client.create_download_task(task["model_id"], task["source"] or "HUGGING_FACE")
if dl.result == "created":
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "WAITING", "self", True, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
started += 1
elif dl.result == "already_downloaded":
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "SUCCESS", "others", False, dl.code, dl.message)
self.repo.mark_task_status(task["id"], "ready_to_submit", dl.code, dl.message)
elif dl.result == "running_elsewhere":
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "RUNNING", "others", False, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
elif dl.result == "download_pool_full":
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(300), increment_attempt=True)
else:
self.repo.mark_task_status(task["id"], "failed", dl.code, dl.message, increment_attempt=True)
return started

118
app/scheduler/loop.py Normal file
View File

@@ -0,0 +1,118 @@
from __future__ import annotations
import logging
from datetime import datetime, timezone
from app.clients.modelhub import ModelHubClient
from app.crawlers.download_success import DownloadSuccessCrawler
from app.crawlers.not_adapted import NotAdaptedCrawler
from app.scheduler.budget import Budget
from app.scheduler.downloader import Downloader
from app.scheduler.planner import Planner
from app.scheduler.submitter import Submitter
from app.settings import Settings
from app.storage.repositories import Repository, future_seconds
LOG = logging.getLogger(__name__)
class StrategyLoop:
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
self.settings = settings
self.repo = repo
self.client = client
self.budget = Budget(repo, settings.max_user_tasks)
self.planner = Planner(repo, self.budget, settings.default_gpu_alias_list)
self.submitter = Submitter(repo, client)
self.downloader = Downloader(repo, client, settings.max_self_downloads, settings.download_poll_interval_seconds)
self.not_adapted_crawler = NotAdaptedCrawler(settings, repo, client)
self.download_success_crawler = DownloadSuccessCrawler(settings, repo, client)
self._last_crawler_refresh = None
def refresh_crawlers_if_due(self) -> dict[str, int]:
now = datetime.now(timezone.utc)
if self._last_crawler_refresh is not None:
elapsed = (now - self._last_crawler_refresh).total_seconds()
if elapsed < self.settings.crawler_refresh_seconds:
return {"not_adapted": 0, "download_success": 0}
self._last_crawler_refresh = now
stats = {"not_adapted": 0, "download_success": 0}
try:
stats["not_adapted"] = self.not_adapted_crawler.refresh()
except Exception:
LOG.exception("not-adapted crawler failed")
try:
stats["download_success"] = self.download_success_crawler.refresh()
except Exception:
LOG.exception("download-success crawler failed")
return stats
def mark_ready_tasks(self, limit: int = 50) -> int:
changed = 0
for task in self.repo.due_tasks(("pending",), limit=limit):
if task["is_bounty"]:
if self.settings.bounty_allow_direct_submit:
self.repo.mark_task_status(task["id"], "ready_to_submit")
changed += 1
else:
status = self.client.get_download_status(task["model_id"])
if status == "SUCCESS":
self.repo.upsert_download(task["model_id"], task["source"], "SUCCESS", "others")
self.repo.mark_task_status(task["id"], "ready_to_submit")
changed += 1
else:
self.repo.mark_task_status(task["id"], "pending", "DOWNLOAD_STATUS", status, future_seconds(300))
elif task["is_promote"]:
absent = self.client.is_model_absent_from_modelhub(task["model_id"])
if not absent:
self.repo.mark_task_status(task["id"], "ready_to_submit")
changed += 1
else:
self.repo.mark_task_status(task["id"], "pending", "NOT_IN_MODELHUB", "model not adapted yet", future_seconds(1800))
elif task["known_downloaded_by_others"]:
status = self.client.get_download_status(task["model_id"])
if status == "SUCCESS":
self.repo.upsert_download(task["model_id"], task["source"], "SUCCESS", "others")
self.repo.mark_task_status(task["id"], "ready_to_submit")
changed += 1
elif task["self_download_allowed"]:
self.repo.mark_task_status(task["id"], "waiting_download", "DOWNLOAD_STATUS", status, future_seconds(300))
else:
self.repo.mark_task_status(task["id"], "pending", "DOWNLOAD_STATUS", status, future_seconds(300))
elif task["self_download_allowed"]:
self.repo.mark_task_status(task["id"], "waiting_download")
else:
self.repo.mark_task_status(task["id"], "skipped", "NO_READY_PATH", "candidate is not bounty/promote/downloaded/self-download")
return changed
def run_once(self) -> dict[str, int]:
crawler_stats = self.refresh_crawlers_if_due()
created = self.planner.materialize_tasks()
polled = self.downloader.poll_active_downloads()
released = self.submitter.release_queue_full_backoff()
ready = self.mark_ready_tasks()
submitted = self.submitter.submit_due()
started_downloads = 0
if submitted == 0 and ready == 0:
started_downloads = self.downloader.maybe_start_self_downloads()
return {
"crawled_not_adapted": crawler_stats["not_adapted"],
"crawled_download_success": crawler_stats["download_success"],
"created_tasks": created,
"polled_downloads": polled,
"released_backoff": released,
"ready_tasks": ready,
"submitted_tasks": submitted,
"started_downloads": started_downloads,
}
def run_forever(self, stop_event):
LOG.info("strategy loop started")
while not stop_event.is_set():
try:
stats = self.run_once()
LOG.info("loop stats: %s", stats)
except Exception:
LOG.exception("strategy loop iteration failed")
stop_event.wait(self.settings.poll_interval_seconds)
LOG.info("strategy loop stopped")

24
app/scheduler/planner.py Normal file
View File

@@ -0,0 +1,24 @@
from app.domain.gpu import expand_gpu_aliases, parse_aliases
from app.domain.model_type import engine_for_model
from app.storage.repositories import Repository
from app.scheduler.budget import Budget
class Planner:
def __init__(self, repo: Repository, budget: Budget, default_aliases: list[str]):
self.repo = repo
self.budget = budget
self.default_aliases = default_aliases
def materialize_tasks(self) -> int:
created = 0
for candidate in self.repo.list_candidates():
if engine_for_model(candidate["model_id"]) != "vllm":
continue
aliases = parse_aliases(candidate["target_gpu_aliases"], self.default_aliases)
for gpu_alias, gpu_type in expand_gpu_aliases(aliases):
if not self.budget.can_create_task():
return created
if self.repo.insert_task_if_absent(candidate["model_id"], gpu_alias, gpu_type, candidate["priority"]):
created += 1
return created

44
app/scheduler/submitter.py Normal file
View File

@@ -0,0 +1,44 @@
from app.clients.modelhub import ModelHubClient
from app.domain.vllm_configs import gen_vllm_config
from app.storage.repositories import Repository, future_seconds
class Submitter:
def __init__(self, repo: Repository, client: ModelHubClient):
self.repo = repo
self.client = client
def submit_due(self, limit: int = 20) -> int:
submitted = 0
for task in self.repo.due_tasks(("ready_to_submit",), limit=limit):
config = gen_vllm_config(task["gpu_type"], task["model_id"], task["max_model_len"])
result = self.client.create_contest_task(task["model_id"], task["gpu_type"], config)
self.repo.record_submission_attempt(
task["id"],
task["model_id"],
task["gpu_type"],
self.client.settings.strategy_id,
self.client.settings.contest_task_strategy_field if self.client.settings.inject_strategy_id else "",
result.code,
result.message,
result.result,
)
if result.result == "success":
self.repo.mark_task_status(task["id"], "submitted", result.code, result.message, increment_attempt=True, increment_submit=True)
submitted += 1
elif result.result == "conflict":
self.repo.mark_task_status(task["id"], "conflict", result.code, result.message, increment_attempt=True)
elif result.result == "queue_full":
backoff = min(3600, 300 * (task["attempt_count"] + 1))
self.repo.mark_task_status(task["id"], "queue_full_backoff", result.code, result.message, future_seconds(backoff), increment_attempt=True)
else:
backoff = min(1800, 120 * (task["attempt_count"] + 1))
self.repo.mark_task_status(task["id"], "pending", result.code, result.message, future_seconds(backoff), increment_attempt=True)
return submitted
def release_queue_full_backoff(self) -> int:
released = 0
for task in self.repo.due_tasks(("queue_full_backoff",), limit=50):
self.repo.mark_task_status(task["id"], "ready_to_submit")
released += 1
return released

165
app/settings.py Normal file
View File

@@ -0,0 +1,165 @@
from dataclasses import dataclass
import json
import os
from pathlib import Path
from typing import Any
def _load_config_file() -> dict[str, Any]:
explicit = os.getenv("STRATEGY_CONFIG_FILE")
candidates = [explicit] if explicit else ["config.local.json", "config.json"]
for item in candidates:
if not item:
continue
path = Path(item)
if path.exists():
with path.open("r", encoding="utf-8") as fh:
return json.load(fh)
return {}
def _value(config: dict[str, Any], name: str, default: Any = None) -> Any:
if name in os.environ:
return os.environ[name]
if name in config:
return config[name]
lower = name.lower()
if lower in config:
return config[lower]
return default
def _bool_value(config: dict[str, Any], name: str, default: bool) -> bool:
raw = _value(config, name, default)
if isinstance(raw, bool):
return raw
if raw is None:
return default
return str(raw).strip().lower() in {"1", "true", "yes", "y", "on"}
def _int_value(config: dict[str, Any], name: str, default: int) -> int:
raw = _value(config, name, default)
if raw is None or raw == "":
return default
return int(raw)
@dataclass(frozen=True)
class Settings:
auth_token: str
hf_token: str
contest_api_token: str
email: str
user_id: str
contributors: str
strategy_id: str
modelhub_url: str = "https://modelhub.org.cn"
modelhub_api_base: str = "https://modelhub.org.cn/api"
modelhub_adminapi_base: str = "https://modelhub.org.cn/adminApi"
db_path: str = "./state.db"
host: str = "0.0.0.0"
port: int = 8080
poll_interval_seconds: int = 60
download_poll_interval_seconds: int = 300
request_timeout_seconds: int = 60
max_self_downloads: int = 8
max_user_tasks: int = 2000
default_gpu_aliases: str = "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8"
default_max_model_len: int = 1024
submit_dry_run: bool = False
inject_strategy_id: bool = True
contest_task_strategy_field: str = "strategyId"
bounty_allow_direct_submit: bool = True
auto_load_seed_path: str = ""
enable_not_adapted_crawler: bool = True
target_machine_names: str = ""
target_task_level: str = "文本生成"
modelhub_page_size: int = 100
crawler_refresh_seconds: int = 3600
crawler_max_pages: int = 0
enable_download_success_crawler: bool = True
download_success_page_size: int = 50
download_success_max_pages: int = 0
config_file_loaded: str = ""
@property
def default_gpu_alias_list(self) -> list[str]:
return [item.strip() for item in self.default_gpu_aliases.split(",") if item.strip()]
def load_settings(require_secrets: bool = True) -> Settings:
config = _load_config_file()
explicit = os.getenv("STRATEGY_CONFIG_FILE")
loaded_path = ""
if explicit and Path(explicit).exists():
loaded_path = explicit
elif Path("config.local.json").exists():
loaded_path = "config.local.json"
elif Path("config.json").exists():
loaded_path = "config.json"
hf_token = _value(config, "HF_TOKEN", None)
if hf_token is None:
hf_token = _value(config, "HFTOKEN", "")
settings = Settings(
auth_token=str(_value(config, "AUTH_TOKEN", "") or ""),
hf_token=str(hf_token or ""),
contest_api_token=str(_value(config, "CONFTEST_API_TOKEN", "") or ""),
email=str(_value(config, "EMAIL", "") or ""),
user_id=str(_value(config, "USER_ID", "") or ""),
contributors=str(_value(config, "CONTRIBUTORS", "") or ""),
strategy_id=str(_value(config, "STRATEGY_ID", "") or ""),
modelhub_url=str(_value(config, "MODELHUB_URL", "https://modelhub.org.cn")),
modelhub_api_base=str(_value(config, "MODELHUB_API_BASE", "https://modelhub.org.cn/api")),
modelhub_adminapi_base=str(_value(config, "MODELHUB_ADMINAPI_BASE", "https://modelhub.org.cn/adminApi")),
db_path=str(_value(config, "DB_PATH", "./state.db")),
host=str(_value(config, "HOST", "0.0.0.0")),
port=_int_value(config, "PORT", 8080),
poll_interval_seconds=_int_value(config, "POLL_INTERVAL_SECONDS", 60),
download_poll_interval_seconds=_int_value(config, "DOWNLOAD_POLL_INTERVAL_SECONDS", 300),
request_timeout_seconds=_int_value(config, "REQUEST_TIMEOUT_SECONDS", 60),
max_self_downloads=_int_value(config, "MAX_SELF_DOWNLOADS", 8),
max_user_tasks=_int_value(config, "MAX_USER_TASKS", 2000),
default_gpu_aliases=str(_value(config, "DEFAULT_GPU_ALIASES", "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8")),
default_max_model_len=_int_value(config, "DEFAULT_MAX_MODEL_LEN", 1024),
submit_dry_run=_bool_value(config, "SUBMIT_DRY_RUN", False),
inject_strategy_id=_bool_value(config, "INJECT_STRATEGY_ID", True),
contest_task_strategy_field=str(_value(config, "CONTEST_TASK_STRATEGY_FIELD", "strategyId")),
bounty_allow_direct_submit=_bool_value(config, "BOUNTY_ALLOW_DIRECT_SUBMIT", True),
auto_load_seed_path=str(_value(config, "AUTO_LOAD_SEED_PATH", "") or ""),
enable_not_adapted_crawler=_bool_value(config, "ENABLE_NOT_ADAPTED_CRAWLER", True),
target_machine_names=str(_value(config, "TARGET_MACHINE_NAMES", "") or ""),
target_task_level=str(_value(config, "TARGET_TASK_LEVEL", "文本生成")),
modelhub_page_size=_int_value(config, "MODELHUB_PAGE_SIZE", 100),
crawler_refresh_seconds=_int_value(config, "CRAWLER_REFRESH_SECONDS", 3600),
crawler_max_pages=_int_value(config, "CRAWLER_MAX_PAGES", 0),
enable_download_success_crawler=_bool_value(config, "ENABLE_DOWNLOAD_SUCCESS_CRAWLER", True),
download_success_page_size=_int_value(config, "DOWNLOAD_SUCCESS_PAGE_SIZE", 50),
download_success_max_pages=_int_value(config, "DOWNLOAD_SUCCESS_MAX_PAGES", 0),
config_file_loaded=loaded_path,
)
if require_secrets:
missing = []
required_names = [
("AUTH_TOKEN", settings.auth_token),
("CONFTEST_API_TOKEN", settings.contest_api_token),
("EMAIL", settings.email),
("USER_ID", settings.user_id),
("CONTRIBUTORS", settings.contributors),
]
for name, value in required_names:
if not value:
missing.append(name)
if settings.inject_strategy_id and not settings.strategy_id:
missing.append("STRATEGY_ID")
if missing:
raise RuntimeError(f"Missing required settings: {', '.join(missing)}")
return settings

14
app/shutdown.py Normal file
View File

@@ -0,0 +1,14 @@
import signal
import threading
stop_event = threading.Event()
def request_shutdown(signum=None, frame=None):
stop_event.set()
def install_signal_handlers():
signal.signal(signal.SIGTERM, request_shutdown)
signal.signal(signal.SIGINT, request_shutdown)

1
app/storage/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""SQLite storage helpers."""

18
app/storage/db.py Normal file
View File

@@ -0,0 +1,18 @@
from __future__ import annotations
import os
import sqlite3
from app.storage.schema import init_schema
def connect(db_path: str) -> sqlite3.Connection:
parent = os.path.dirname(os.path.abspath(db_path))
if parent and not os.path.exists(parent):
os.makedirs(parent, exist_ok=True)
conn = sqlite3.connect(db_path, timeout=30)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=30000")
init_schema(conn)
return conn

266
app/storage/repositories.py Normal file
View File

@@ -0,0 +1,266 @@
from __future__ import annotations
from datetime import datetime, timezone, timedelta
import sqlite3
def utcnow() -> str:
return datetime.now(timezone.utc).isoformat()
def future_seconds(seconds: int) -> str:
return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat()
class Repository:
def __init__(self, conn: sqlite3.Connection):
self.conn = conn
def upsert_candidate(
self,
model_id: str,
source: str,
origin: str,
priority: int,
is_bounty: bool = False,
is_promote: bool = False,
known_downloaded_by_others: bool = False,
self_download_allowed: bool = False,
max_model_len: int = 1024,
target_gpu_aliases: str | None = None,
notes: str | None = None,
) -> None:
now = utcnow()
existing = self.conn.execute(
"SELECT target_gpu_aliases FROM candidates WHERE model_id = ?",
(model_id,),
).fetchone()
if existing and target_gpu_aliases:
old_aliases = [item.strip() for item in (existing["target_gpu_aliases"] or "").split(",") if item.strip()]
new_aliases = [item.strip() for item in target_gpu_aliases.split(",") if item.strip()]
merged = []
for alias in old_aliases + new_aliases:
if alias not in merged:
merged.append(alias)
target_gpu_aliases = ",".join(merged)
self.conn.execute(
"""
INSERT INTO candidates (
model_id, source, origin, priority, is_bounty, is_promote,
known_downloaded_by_others, self_download_allowed, max_model_len,
target_gpu_aliases, notes, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(model_id) DO UPDATE SET
source=excluded.source,
origin=excluded.origin,
priority=min(candidates.priority, excluded.priority),
is_bounty=max(candidates.is_bounty, excluded.is_bounty),
is_promote=max(candidates.is_promote, excluded.is_promote),
known_downloaded_by_others=max(candidates.known_downloaded_by_others, excluded.known_downloaded_by_others),
self_download_allowed=max(candidates.self_download_allowed, excluded.self_download_allowed),
max_model_len=excluded.max_model_len,
target_gpu_aliases=excluded.target_gpu_aliases,
notes=excluded.notes,
updated_at=excluded.updated_at
""",
(
model_id,
source,
origin,
priority,
int(is_bounty),
int(is_promote),
int(known_downloaded_by_others),
int(self_download_allowed),
max_model_len,
target_gpu_aliases,
notes,
now,
now,
),
)
self.conn.commit()
def list_candidates(self) -> list[sqlite3.Row]:
cursor = self.conn.execute("SELECT * FROM candidates ORDER BY priority, created_at, model_id")
return list(cursor.fetchall())
def insert_task_if_absent(self, model_id: str, gpu_alias: str, gpu_type: str, priority: int) -> bool:
now = utcnow()
cur = self.conn.execute(
"""
INSERT OR IGNORE INTO tasks (
model_id, gpu_alias, gpu_type, engine, priority, status,
created_at, updated_at
) VALUES (?, ?, ?, 'vllm', ?, 'pending', ?, ?)
""",
(model_id, gpu_alias, gpu_type, priority, now, now),
)
self.conn.commit()
return cur.rowcount > 0
def count_budgeted_tasks(self) -> int:
cursor = self.conn.execute("SELECT COUNT(*) FROM tasks WHERE status != 'skipped'")
return int(cursor.fetchone()[0])
def mark_task_status(
self,
task_id: int,
status: str,
error_code: str | None = None,
error_message: str | None = None,
next_attempt_at: str | None = None,
increment_attempt: bool = False,
increment_submit: bool = False,
) -> None:
now = utcnow()
self.conn.execute(
"""
UPDATE tasks SET
status = ?,
last_error_code = ?,
last_error_message = ?,
next_attempt_at = ?,
attempt_count = attempt_count + ?,
submit_count = submit_count + ?,
submitted_at = CASE WHEN ? = 'submitted' THEN ? ELSE submitted_at END,
completed_at = CASE WHEN ? IN ('submitted', 'conflict', 'failed', 'skipped') THEN ? ELSE completed_at END,
updated_at = ?
WHERE id = ?
""",
(
status,
error_code,
error_message,
next_attempt_at,
1 if increment_attempt else 0,
1 if increment_submit else 0,
status,
now,
status,
now,
now,
task_id,
),
)
self.conn.commit()
def due_tasks(self, statuses: tuple[str, ...], limit: int = 50) -> list[sqlite3.Row]:
placeholders = ",".join("?" for _ in statuses)
now = utcnow()
cursor = self.conn.execute(
f"""
SELECT t.*, c.is_bounty, c.is_promote, c.known_downloaded_by_others,
c.self_download_allowed, c.max_model_len, c.source
FROM tasks t
JOIN candidates c ON c.model_id = t.model_id
WHERE t.status IN ({placeholders})
AND (t.next_attempt_at IS NULL OR t.next_attempt_at <= ?)
ORDER BY t.priority, t.attempt_count, t.created_at, t.model_id, t.gpu_type
LIMIT ?
""",
(*statuses, now, limit),
)
return list(cursor.fetchall())
def count_active_self_downloads(self) -> int:
cursor = self.conn.execute(
"SELECT COUNT(*) FROM downloads WHERE active_self_slot = 1 AND status IN ('WAITING', 'RUNNING', 'unknown')"
)
return int(cursor.fetchone()[0])
def upsert_download(
self,
model_id: str,
source: str,
status: str,
owner: str,
active_self_slot: bool = False,
download_code: str | None = None,
download_message: str | None = None,
next_poll_at: str | None = None,
) -> None:
now = utcnow()
self.conn.execute(
"""
INSERT INTO downloads (
model_id, source, status, owner, active_self_slot, download_code,
download_message, last_polled_at, next_poll_at, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(model_id) DO UPDATE SET
source=excluded.source,
status=excluded.status,
owner=excluded.owner,
active_self_slot=excluded.active_self_slot,
download_code=excluded.download_code,
download_message=excluded.download_message,
last_polled_at=excluded.last_polled_at,
next_poll_at=excluded.next_poll_at,
updated_at=excluded.updated_at
""",
(
model_id,
source,
status,
owner,
int(active_self_slot),
download_code,
download_message,
now,
next_poll_at,
now,
now,
),
)
self.conn.commit()
def get_download(self, model_id: str):
cursor = self.conn.execute("SELECT * FROM downloads WHERE model_id = ?", (model_id,))
return cursor.fetchone()
def active_downloads_due(self, limit: int = 20) -> list[sqlite3.Row]:
now = utcnow()
cursor = self.conn.execute(
"""
SELECT * FROM downloads
WHERE active_self_slot = 1
AND status IN ('WAITING', 'RUNNING', 'unknown')
AND (next_poll_at IS NULL OR next_poll_at <= ?)
ORDER BY next_poll_at, created_at
LIMIT ?
""",
(now, limit),
)
return list(cursor.fetchall())
def record_submission_attempt(
self,
task_id: int,
model_id: str,
gpu_type: str,
strategy_id: str,
payload_strategy_field: str,
response_code: str | None,
response_message: str | None,
result: str,
) -> None:
self.conn.execute(
"""
INSERT INTO submission_attempts (
task_id, model_id, gpu_type, strategy_id, payload_strategy_field,
response_code, response_message, result, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
task_id,
model_id,
gpu_type,
strategy_id,
payload_strategy_field,
response_code,
response_message,
result,
utcnow(),
),
)
self.conn.commit()

87
app/storage/schema.py Normal file
View File

@@ -0,0 +1,87 @@
SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS candidates (
model_id TEXT PRIMARY KEY,
source TEXT NOT NULL DEFAULT 'HUGGING_FACE',
origin TEXT NOT NULL,
priority INTEGER NOT NULL,
is_bounty INTEGER NOT NULL DEFAULT 0,
is_promote INTEGER NOT NULL DEFAULT 0,
known_downloaded_by_others INTEGER NOT NULL DEFAULT 0,
self_download_allowed INTEGER NOT NULL DEFAULT 0,
max_model_len INTEGER NOT NULL DEFAULT 1024,
target_gpu_aliases TEXT,
notes TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS tasks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
model_id TEXT NOT NULL,
gpu_alias TEXT NOT NULL,
gpu_type TEXT NOT NULL,
engine TEXT NOT NULL DEFAULT 'vllm',
priority INTEGER NOT NULL,
status TEXT NOT NULL DEFAULT 'pending',
attempt_count INTEGER NOT NULL DEFAULT 0,
submit_count INTEGER NOT NULL DEFAULT 0,
last_error_code TEXT,
last_error_message TEXT,
next_attempt_at TEXT,
submitted_at TEXT,
completed_at TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
UNIQUE(model_id, gpu_type, engine)
);
CREATE TABLE IF NOT EXISTS downloads (
model_id TEXT PRIMARY KEY,
source TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'unknown',
owner TEXT NOT NULL DEFAULT 'unknown',
active_self_slot INTEGER NOT NULL DEFAULT 0,
download_code TEXT,
download_message TEXT,
waiting_since TEXT,
last_polled_at TEXT,
next_poll_at TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS submission_attempts (
id INTEGER PRIMARY KEY AUTOINCREMENT,
task_id INTEGER NOT NULL,
model_id TEXT NOT NULL,
gpu_type TEXT NOT NULL,
strategy_id TEXT,
payload_strategy_field TEXT,
response_code TEXT,
response_message TEXT,
result TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS seed_imports (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_path TEXT NOT NULL,
file_hash TEXT,
imported_count INTEGER NOT NULL,
created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_tasks_sched
ON tasks(status, priority, next_attempt_at, created_at);
CREATE INDEX IF NOT EXISTS idx_tasks_model
ON tasks(model_id);
CREATE INDEX IF NOT EXISTS idx_downloads_poll
ON downloads(status, next_poll_at);
"""
def init_schema(conn):
conn.executescript(SCHEMA_SQL)
conn.commit()