modelhub-submit-agent/main.py

"""
ModelHub Cancel-All 智能体
启动后自动取消当前账号所有 waiting/running 状态的验证任务
用法：推送此版本到平台，选择对应 tag 运行即可
"""

import json
import os
import signal
import time
import threading
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import requests

# ============================================================
# 配置
# ============================================================

HOST = "0.0.0.0"
PORT = 8080

# 目标GPU和Token（跟提交版本保持一致）
TARGET_GPU = "Iluvatar_bi-150"
TARGET_TOKEN = "f45f1aae2c094426be237c88b1085015"

MODELHUB_API = "https://modelhub.org.cn/api"
BATCH_SIZE = 50  # 每批取消数量（API上限）

# ============================================================
# 全局状态
# ============================================================

state = {'running': False, 'last_run': None, 'logs': []}
state_lock = threading.Lock()


def log(msg: str):
    ts = datetime.now().strftime('%H:%M:%S')
    line = f"[{ts}] {msg}"
    print(line, flush=True)
    with state_lock:
        state['logs'].append(line)
        if len(state['logs']) > 500:
            state['logs'] = state['logs'][-300:]


def cancel_all_tasks():
    """取消所有 waiting 和 running 状态的任务"""
    headers = {
        'Xc-Token': TARGET_TOKEN,
        'Accept': 'application/json',
        'Content-Type': 'application/json',
    }

    log("=" * 50)
    log("开始取消所有任务")
    log(f"目标GPU: {TARGET_GPU}")

    # 1. 查询所有任务
    all_tasks = []
    page = 1
    while True:
        try:
            resp = requests.get(
                f"{MODELHUB_API}/adapt/task/page",
                headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'},
                params={'current': page, 'pageSize': 100, 'onlyMine': 'true',
                        'gpuType': TARGET_GPU},
                timeout=15
            )
            data = resp.json()
            if data.get('code') != 0:
                log(f"查询失败: code={data.get('code')}")
                break
            records = data['data'].get('records', [])
            if not records:
                break
            all_tasks.extend(records)
            total = data['data'].get('total', 0)
            log(f"  第{page}页: {len(records)} 条, 累计 {len(all_tasks)} / {total}")
            if len(all_tasks) >= int(total):
                break
            page += 1
        except Exception as e:
            log(f"  查询异常: {e}")
            return

    # 2. 筛选 waiting/running 状态
    cancellable = [t for t in all_tasks if t.get('status') in ('waiting', 'running')]
    if not cancellable:
        log(f"  没有可取消的任务 (总任务 {len(all_tasks)} 个)")
        log("=" * 50)
        return

    task_ids = [int(t.get('taskId')) for t in cancellable if t.get('taskId')]
    log(f"  总任务 {len(all_tasks)} 个, 可取消 {len(task_ids)} 个")

    # 3. 分批取消
    url = f"{MODELHUB_API}/async/task/stop-create-contest-task"
    cancelled = 0
    for i in range(0, len(task_ids), BATCH_SIZE):
        batch = task_ids[i:i + BATCH_SIZE]
        try:
            resp = requests.put(url, headers=headers, json={'taskIds': batch}, timeout=15)
            data = resp.json()
            if data.get('code') == 0:
                cancelled += len(batch)
                log(f"  批次 {i // BATCH_SIZE + 1}: 取消 {len(batch)} 个 OK")
            else:
                log(f"  批次 {i // BATCH_SIZE + 1}: 失败 code={data.get('code')} msg={data.get('message', '')[:80]}")
        except Exception as e:
            log(f"  批次 {i // BATCH_SIZE + 1}: 异常 {e}")
        time.sleep(1)

    log(f"  取消完成: {cancelled} / {len(task_ids)}")
    log("=" * 50)


# ============================================================
# HTTP 服务
# ============================================================

class AgentHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/health':
            self._json({'status': 'ok'})
        elif self.path == '/':
            self._json({
                'name': 'modelhub-cancel-all',
                'gpu': TARGET_GPU,
                'status': 'running' if state['running'] else 'idle',
                'last_result': state.get('last_result'),
            })
        elif self.path.startswith('/logs'):
            lines = 100
            self._json({'logs': state['logs'][-lines:]})
        else:
            self._json({'error': 'not found'}, 404)

    def _json(self, body: dict, status: int = 200):
        payload = json.dumps(body, ensure_ascii=False).encode()
        self.send_response(status)
        self.send_header('Content-Type', 'application/json; charset=utf-8')
        self.send_header('Content-Length', str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def log_message(self, fmt, *args):
        pass


# ============================================================
# 启动
# ============================================================

shutdown_requested = False


def _handle_signal(signum, _frame):
    global shutdown_requested
    shutdown_requested = True
    log(f"收到信号 {signum}，准备关闭")


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

    server = ThreadingHTTPServer((HOST, PORT), AgentHandler)
    server.timeout = 1

    log(f"Cancel-All 智能体启动 | {HOST}:{PORT}")
    log(f"目标GPU: {TARGET_GPU}")

    # 启动后自动取消所有任务
    def _auto_cancel():
        time.sleep(2)
        try:
            state['running'] = True
            state['last_run'] = datetime.now().isoformat()
            cancel_all_tasks()
            state['last_result'] = {'success': True}
        except Exception as e:
            log(f"异常: {e}")
            state['last_result'] = {'error': str(e)}
        finally:
            state['running'] = False

    threading.Thread(target=_auto_cancel, daemon=True).start()

    while not shutdown_requested:
        server.handle_request()

    server.server_close()
    log("智能体已关闭")


if __name__ == '__main__':
    main()