add main.py
This commit is contained in:
239
main.py
Normal file
239
main.py
Normal file
@@ -0,0 +1,239 @@
|
||||
"""ModelHub XC 适配智能体 —— 单任务提交器
|
||||
|
||||
启动后向 ModelHub XC 提交 1 条模型适配验证任务，随后保持 /health 存活等待平台停止。
|
||||
|
||||
提交内容（已查重，非重复模型）：
|
||||
模型：https://www.modelscope.cn/models/qnguyen3/qwq-1.5b-grpo (Qwen2 / 1.5B)
|
||||
算力：Ascend_910-b4 框架：vllm 任务类型：text-generation
|
||||
|
||||
平台契约：
|
||||
- 根目录 Dockerfile，监听 8080
|
||||
- GET /health 探活
|
||||
- 读取平台注入的 STRATEGY_ID
|
||||
- 正确处理 SIGTERM
|
||||
"""
|
||||
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", "")
|
||||
|
||||
MODEL_ADDRESS = "https://www.modelscope.cn/models/qnguyen3/qwq-1.5b-grpo"
|
||||
TASK_TYPE = "text-generation"
|
||||
TARGET_GPU = "Ascend_910-b4"
|
||||
FRAMEWORK = "vllm"
|
||||
|
||||
# 与 POST /api/adapt/task/build-config?gpuType=Ascend_910-b4&framework=vllm&taskType=text-generation 输出一致
|
||||
CONFIG_PARAMS = """framework: vllm
|
||||
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
|
||||
"""
|
||||
|
||||
HEADERS = {"Content-Type": "application/json", "Xc-Token": XC_TOKEN}
|
||||
|
||||
HTTP_HOST, HTTP_PORT = "0.0.0.0", 8080
|
||||
MAX_ATTEMPTS = 5
|
||||
RETRY_GAP = 15 # 秒
|
||||
|
||||
_state = {
|
||||
"strategy_id": STRATEGY_ID,
|
||||
"phase": "starting", # starting | submitting | done | error
|
||||
"model_address": MODEL_ADDRESS,
|
||||
"target_gpu": TARGET_GPU,
|
||||
"framework": FRAMEWORK,
|
||||
"attempts": 0,
|
||||
"attempt_log": [],
|
||||
"task_id": None,
|
||||
"started_at": None,
|
||||
"finished_at": None,
|
||||
}
|
||||
_shutdown = threading.Event()
|
||||
_submitting = threading.Lock()
|
||||
_submitted = False
|
||||
|
||||
|
||||
def submit_task():
|
||||
"""提交一次；已提交过则直接返回。"""
|
||||
global _submitted
|
||||
with _submitting:
|
||||
if _submitted:
|
||||
return True
|
||||
|
||||
payload = {
|
||||
"configParams": CONFIG_PARAMS,
|
||||
"framework": FRAMEWORK,
|
||||
"modelAddress": MODEL_ADDRESS,
|
||||
"targetGpu": TARGET_GPU,
|
||||
"taskType": TASK_TYPE,
|
||||
"strategyId": STRATEGY_ID,
|
||||
}
|
||||
|
||||
_state["phase"] = "submitting"
|
||||
_state["started_at"] = datetime.utcnow().isoformat()
|
||||
|
||||
for i in range(1, MAX_ATTEMPTS + 1):
|
||||
if _shutdown.is_set():
|
||||
break
|
||||
_state["attempts"] = i
|
||||
try:
|
||||
resp = requests.post(BASE_URL + ADD_TASK_ENDPOINT, headers=HEADERS,
|
||||
json=payload, timeout=60)
|
||||
body = resp.json()
|
||||
except Exception as e: # noqa: BLE001
|
||||
body = {"code": -1, "message": "exception: %s" % e}
|
||||
|
||||
code = body.get("code")
|
||||
msg = body.get("message", "")
|
||||
entry = {"attempt": i, "time": datetime.utcnow().isoformat(),
|
||||
"http_status": None, "code": code, "message": msg}
|
||||
try:
|
||||
entry["http_status"] = resp.status_code
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
_state["attempt_log"].append(entry)
|
||||
print("[submit] #%d code=%s message=%s" % (i, code, msg), flush=True)
|
||||
|
||||
if code == 0:
|
||||
data = body.get("data") or {}
|
||||
_state["task_id"] = data.get("id")
|
||||
_state["phase"] = "done"
|
||||
_state["finished_at"] = datetime.utcnow().isoformat()
|
||||
_submitted = True
|
||||
print("[submit] 提交成功 taskId=%s" % _state["task_id"], flush=True)
|
||||
return True
|
||||
|
||||
# 已存在/重复等不可重试错误，直接退出
|
||||
if code not in (-1, 60014, 50000):
|
||||
_state["phase"] = "error"
|
||||
_state["finished_at"] = datetime.utcnow().isoformat()
|
||||
print("[submit] 终止：不可重试错误", flush=True)
|
||||
return False
|
||||
|
||||
if i < MAX_ATTEMPTS and not _shutdown.is_set():
|
||||
time.sleep(RETRY_GAP)
|
||||
|
||||
_state["phase"] = "error"
|
||||
_state["finished_at"] = datetime.utcnow().isoformat()
|
||||
return False
|
||||
|
||||
|
||||
class Handler(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 do_POST(self):
|
||||
if self.path == "/task":
|
||||
submit_task()
|
||||
self._json(_state)
|
||||
else:
|
||||
self._json({"error": "not found"}, 404)
|
||||
|
||||
def _json(self, body, code=200):
|
||||
raw = json.dumps(body, ensure_ascii=False, default=str).encode()
|
||||
self.send_response(code)
|
||||
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, fmt, *args):
|
||||
print(self.address_string() + " - " + (fmt % args), flush=True)
|
||||
|
||||
|
||||
def _handle_signal(signum, _frame):
|
||||
print("[main] 收到信号 %s，准备退出" % signum, flush=True)
|
||||
_shutdown.set()
|
||||
|
||||
|
||||
def main():
|
||||
signal.signal(signal.SIGTERM, _handle_signal)
|
||||
signal.signal(signal.SIGINT, _handle_signal)
|
||||
|
||||
print("[main] strategy_id=%s" % STRATEGY_ID, flush=True)
|
||||
print("[main] model=%s gpu=%s framework=%s" % (MODEL_ADDRESS, TARGET_GPU, FRAMEWORK), flush=True)
|
||||
|
||||
http_thread = threading.Thread(target=_run_http, daemon=False)
|
||||
http_thread.start()
|
||||
|
||||
worker = threading.Thread(target=submit_task, daemon=True)
|
||||
worker.start()
|
||||
|
||||
_shutdown.wait()
|
||||
print("[main] 退出", flush=True)
|
||||
|
||||
|
||||
def _run_http():
|
||||
server = ThreadingHTTPServer((HTTP_HOST, HTTP_PORT), Handler)
|
||||
server.timeout = 1
|
||||
print("[http] 监听 %s:%d" % (HTTP_HOST, HTTP_PORT), flush=True)
|
||||
while not _shutdown.is_set():
|
||||
server.handle_request()
|
||||
server.server_close()
|
||||
print("[http] 已关闭", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user