109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
"""ModelHub XC 适配智能体 —— 低排队算力批次（5 条）
|
||
|
||
每条任务分配一张低排队算力卡，并选用该卡历史通过率较高的框架。
|
||
OSP 系列 max_position_embeddings=2048，配置里的 max_model_len 已由 4096 下调为 2048。
|
||
"""
|
||
import json
|
||
import os
|
||
import signal
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
import requests
|
||
|
||
BASE_URL = os.environ.get("BASE_URL", "https://modelhub.org.cn")
|
||
ADD_TASK_ENDPOINT = "/api/adapt/task/add"
|
||
XC_TOKEN = os.environ.get("XC_TOKEN", "6616dfe75c874e61987b1f7c908e6d77")
|
||
STRATEGY_ID = os.environ.get("STRATEGY_ID", "")
|
||
TASK_TYPE = "text-generation"
|
||
|
||
TASKS = [
|
||
{
|
||
"label": "OSP-1.4B-100B-Shampoo-SSNorm-EmbProj",
|
||
"model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-100B-Shampoo-SSNorm-EmbProj",
|
||
"gpu": "Kunlunxin_p-800",
|
||
"framework": "vllm_fix_tokenizer",
|
||
"config": "framework: vllm_fix_tokenizer\nlang: en\nmax_model_len: 2048\nsut_config:\n gpu_num: 1\n values:\n command:\n - /opt/entrypoint.sh - /model\n - --port\n - '80'\n - --served-model-name\n - llm\n - --max-model-len\n - '2048'\n - --enforce-eager\n - --trust-remote-code\n - -tp\n - '1'\nref_config:\n gpu_num: 1\n values:\n command:\n - /opt/entrypoint.sh - /model\n - --port\n - '80'\n - --served-model-name\n - llm\n - --max-model-len\n - '2048'\n - --enforce-eager\n - --trust-remote-code\n - -tp\n - '1'"
|
||
}
|
||
]
|
||
|
||
HEADERS = {"Content-Type": "application/json", "Xc-Token": XC_TOKEN}
|
||
HTTP_HOST, HTTP_PORT = "0.0.0.0", 8080
|
||
GAP = 4
|
||
_shutdown = threading.Event()
|
||
_state = {"phase": "starting", "total": len(TASKS), "submitted": 0, "failed": 0, "results": []}
|
||
|
||
|
||
def submit_one(t):
|
||
body = {
|
||
"configParams": t["config"],
|
||
"framework": t["framework"],
|
||
"modelAddress": t["model"],
|
||
"targetGpu": t["gpu"],
|
||
"taskType": TASK_TYPE,
|
||
"strategyId": STRATEGY_ID,
|
||
}
|
||
try:
|
||
r = requests.post(BASE_URL + ADD_TASK_ENDPOINT, headers=HEADERS, json=body, timeout=90)
|
||
j = r.json()
|
||
except Exception as e: # noqa: BLE001
|
||
j = {"code": -1, "message": "exception: %s" % e}
|
||
tid = (j.get("data") or {}).get("id")
|
||
print("[submit] %s gpu=%s fw=%s code=%s msg=%s taskId=%s"
|
||
% (t["label"], t["gpu"], t["framework"], j.get("code"), j.get("message"), tid), flush=True)
|
||
return j.get("code"), tid
|
||
|
||
|
||
def run():
|
||
_state["phase"] = "submitting"
|
||
for i, t in enumerate(TASKS, 1):
|
||
if _shutdown.is_set():
|
||
break
|
||
code, tid = submit_one(t)
|
||
_state["results"].append({"label": t["label"], "gpu": t["gpu"],
|
||
"framework": t["framework"], "code": code, "task_id": tid})
|
||
_state["submitted" if code == 0 else "failed"] += 1
|
||
if i < len(TASKS):
|
||
time.sleep(GAP)
|
||
_state["phase"] = "done"
|
||
print("[batch] 完成 submitted=%d failed=%d" % (_state["submitted"], _state["failed"]), flush=True)
|
||
|
||
|
||
class H(BaseHTTPRequestHandler):
|
||
def do_GET(self):
|
||
if self.path == "/health":
|
||
self._json({"status": "ok"})
|
||
elif self.path in ("/", "/status"):
|
||
self._json(_state)
|
||
else:
|
||
self._json({"error": "not found"}, 404)
|
||
|
||
def _json(self, b, c=200):
|
||
raw = json.dumps(b, ensure_ascii=False, default=str).encode()
|
||
self.send_response(c)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(raw)))
|
||
self.end_headers()
|
||
self.wfile.write(raw)
|
||
|
||
def log_message(self, *a):
|
||
pass
|
||
|
||
|
||
def main():
|
||
signal.signal(signal.SIGTERM, lambda *_: _shutdown.set())
|
||
print("[main] strategy_id=%s tasks=%d" % (STRATEGY_ID, len(TASKS)), flush=True)
|
||
for t in TASKS:
|
||
print("[main] 计划: %s @ %s / %s" % (t["label"], t["gpu"], t["framework"]), flush=True)
|
||
srv = ThreadingHTTPServer((HTTP_HOST, HTTP_PORT), H)
|
||
srv.timeout = 1
|
||
threading.Thread(target=run, daemon=True).start()
|
||
while not _shutdown.is_set():
|
||
srv.handle_request()
|
||
srv.server_close()
|
||
|
||
|
||
main()
|