xc-adapt-agent-2174/main.py

"""ModelHub XC 适配智能体 —— 多任务提交器（本批次 4 条）

启动后依次提交 4 条已查重的模型适配验证任务，随后保持 /health 存活等待平台停止。

批次清单（提交前均查过：模型列表 0 条 / 全站任务 0 条 / 详情接口无验证记录）：
  1. mrm8488/DeepSeek-R1-Distill-Qwen-1.5B-grpo-limo-2e   Qwen2 1.5B  Ascend_910-b4  vllm
  2. HollowInky/skill-search-qwen3-4b-unmasked-grpo       Qwen3 4B    Ascend_910-b4  vllm
  3. dmis-lab/OSP-1.4B-100B-Muon                          Llama 1.4B  Ascend_910-b4  vllm
  4. qnguyen3/quan-1.8b-chat                              Llama 1.8B  Ascend_910-b3  vllm

平台契约：根目录 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", "")

# 910B4 模板（max_model_len 2048）
CFG_B4 = """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"""

# 910B3 模板（max_model_len 4096）
CFG_B3 = """framework: vllm
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
    - '20644'
    - --served-model-name
    - llm
    - --max-model-len
    - '4096'
    - -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"""

TASKS = [
    {
        "label": "OSP-1.4B-100B-Adam",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-100B-Adam",
        "gpu": "Ascend_910-b4",
        "config": CFG_B4,
    },
    {
        "label": "OSP-1.4B-1T-Adam",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-1T-Adam",
        "gpu": "Ascend_910-b4",
        "config": CFG_B4,
    },
]

TASK_TYPE = "text-generation"
FRAMEWORK = "vllm"
HEADERS = {"Content-Type": "application/json", "Xc-Token": XC_TOKEN}

HTTP_HOST, HTTP_PORT = "0.0.0.0", 8080
GAP = 4  # 两条之间的间隔秒数

_state = {
    "strategy_id": STRATEGY_ID,
    "phase": "starting",   # starting | submitting | done
    "total": len(TASKS),
    "submitted": 0,
    "failed": 0,
    "results": [],
    "started_at": None,
    "finished_at": None,
}
_shutdown = threading.Event()


def submit_one(task):
    payload = {
        "configParams": task["config"],
        "framework": FRAMEWORK,
        "modelAddress": task["model"],
        "targetGpu": task["gpu"],
        "taskType": TASK_TYPE,
        "strategyId": STRATEGY_ID,
    }
    try:
        resp = requests.post(BASE_URL + ADD_TASK_ENDPOINT, headers=HEADERS,
                             json=payload, timeout=90)
        body = resp.json()
    except Exception as e:  # noqa: BLE001
        body = {"code": -1, "message": "exception: %s" % e}

    code = body.get("code")
    data = body.get("data") or {}
    print("[submit] %s gpu=%s code=%s message=%s taskId=%s"
          % (task["label"], task["gpu"], code, body.get("message"), data.get("id")), flush=True)
    return code, body.get("message"), data.get("id")


def run_batch():
    _state["phase"] = "submitting"
    _state["started_at"] = datetime.utcnow().isoformat()

    for i, task in enumerate(TASKS, 1):
        if _shutdown.is_set():
            break
        code, msg, task_id = submit_one(task)
        _state["results"].append({
            "index": i,
            "label": task["label"],
            "model": task["model"],
            "gpu": task["gpu"],
            "code": code,
            "message": msg,
            "task_id": task_id,
            "time": datetime.utcnow().isoformat(),
        })
        if code == 0:
            _state["submitted"] += 1
        else:
            _state["failed"] += 1
        if i < len(TASKS):
            time.sleep(GAP)

    _state["phase"] = "done"
    _state["finished_at"] = datetime.utcnow().isoformat()
    print("[batch] 完成 submitted=%d failed=%d total=%d"
          % (_state["submitted"], _state["failed"], _state["total"]), flush=True)


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 _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 _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()


def main():
    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    print("[main] strategy_id=%s tasks=%d" % (STRATEGY_ID, len(TASKS)), flush=True)
    for t in TASKS:
        print("[main] 计划: %s @ %s" % (t["label"], t["gpu"]), flush=True)

    http_thread = threading.Thread(target=_run_http, daemon=False)
    http_thread.start()

    worker = threading.Thread(target=run_batch, daemon=True)
    worker.start()

    _shutdown.wait()


if __name__ == "__main__":
    main()