Stop Sunrise waiting tasks for three accounts

main.py

@@ -1,14 +1,15 @@
-"""Stop waiting ModelHub XC validation tasks for the configured accounts.
+"""Stop waiting Sunrise_pt-200-x1 validation tasks for selected accounts.
 
-The service performs the stop requests once at startup and then stays alive so
-the strategy platform can probe it through ``/health`` and inspect ``/status``.
+The worker processes accounts sequentially. For each account it discovers the
+current target tasks, validates account/status/GPU locally, then stops them in
+batches with that account's own xc-Token. The HTTP service remains available
+for platform health and status probes after the worker finishes.
 """
 
 import json
 import os
 import signal
 import threading
-import time
 from datetime import datetime, timezone
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import Any
@@ -16,86 +17,43 @@ from typing import Any
 import requests
 
 BASE_URL = os.environ.get("BASE_URL", "https://modelhub.org.cn").rstrip("/")
+TASK_PAGE_ENDPOINT = "/api/adapt/task/page"
 STOP_TASK_ENDPOINT = "/api/async/task/stop-create-contest-task"
+TARGET_GPU = "Sunrise_pt-200-x1"
+TARGET_STATUS = "waiting"
 STRATEGY_ID = os.environ.get("STRATEGY_ID", "")
 
 HTTP_HOST = "0.0.0.0"
 HTTP_PORT = int(os.environ.get("PORT", "8080"))
+PAGE_SIZE = 100
 BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "50"))
 MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "3"))
 REQUEST_TIMEOUT = int(os.environ.get("REQUEST_TIMEOUT", "30"))
 
-# Snapshot queried from /api/adapt/task/page on 2026-07-26. Only records
-# whose status was exactly "waiting" are included.
 TARGETS = [
     {
-        "account": "fanyi",
-        "xc_token": "f2d501c9ae6543a589cd6cb789108c41",
-        "task_ids": [
-            3469928, 3469927, 3469926, 3469925, 3469924, 3469923, 3469922,
-            3469921, 3469920, 3469919, 3469918, 3469917, 3469916, 3469915,
-            3469914, 3469913, 3469911, 3469910, 3469909, 3469908, 3469907,
-            3469905, 3469904, 3469903, 3469902, 3469901, 3469900, 3469899,
-            3469898, 3469897, 3469896, 3469895, 3469894, 3469893, 3469892,
-            3469891, 3469890, 3469889, 3469888, 3469887, 3469886, 3469885,
-            3469884, 3469883, 3469882, 3469881, 3469879, 3469878, 3469877,
-            3469875, 3469874, 3469873, 3469872, 3469871, 3469870, 3469869,
-            3469868, 3469867, 3469866, 3469865, 3469864, 3469863, 3469862,
-            3469861, 3469860, 3469859, 3469858, 3469857, 3469856, 3469855,
-            3469854, 3469853, 3469852, 3469851, 3469850, 3469849, 3469848,
-            3469847, 3469846, 3469845, 3469844, 3469843, 3469842, 3469841,
-            3469840, 3469837, 3469836, 3469835, 3469834, 3469833, 3469832,
-            3469831, 3469830, 3469829, 3469828, 3469827, 3469826, 3469825,
-            3469822, 3469821,
-        ],
+        "account": "jiangxiaowen",
+        "xc_token": "88d5fee9f1fe4f7583f11a9d3702dc85",
     },
     {
-        "account": "jiajing",
-        "xc_token": "5e051e0ff8384a81af53bea780deb28a",
-        "task_ids": [
-            3469551, 3469550, 3469549, 3469548, 3469547, 3469546, 3469545,
-            3469544, 3469543, 3469542, 3469541, 3469540, 3469539, 3469538,
-            3469537, 3469536, 3469535, 3469534, 3469533, 3469532, 3469531,
-            3469530, 3469529, 3469528, 3469526, 3469525, 3469524, 3469523,
-            3469522, 3469521, 3469520, 3469519, 3469518, 3469517, 3469516,
-            3469515, 3469514, 3469513, 3469512, 3469511, 3469510, 3469508,
-            3469507, 3469503, 3469501, 3469500, 3469499, 3469497, 3469495,
-            3469493, 3469492, 3469491, 3469490, 3469488, 3469486, 3469485,
-            3469484, 3469483, 3469482, 3469480, 3469477, 3469476, 3469475,
-            3469474, 3469472, 3469471, 3469470, 3469469, 3469468, 3469467,
-            3469466, 3469465, 3469464, 3469463, 3469462, 3469461, 3469460,
-            3469459, 3469458, 3469457, 3469456, 3469455, 3469454, 3469453,
-            3469452, 3469451, 3469450, 3469449,
-        ],
+        "account": "l112233",
+        "xc_token": "40cb6910dc9a442a816298a228da65ac",
     },
     {
-        "account": "i-zhouyuanxi@4paradigm.com",
-        "xc_token": "62b9b487eff2488fb9f1da0b963f0b93",
-        "task_ids": [
-            3450582, 3450581, 3450580, 3450579, 3450575, 3450573, 3450572,
-            3450570, 3450569, 3450568, 3450567, 3450565, 3450560, 3450557,
-            3450556, 3450555, 3450554, 3450553, 3450552, 3450551, 3450550,
-            3450549, 3450548, 3450544, 3450541, 3450538, 3450537, 3450536,
-            3450535, 3450534, 3450533, 3450532, 3450531, 3450530, 3450529,
-            3450528, 3450527, 3450526, 3450525, 3450524, 3450523, 3450522,
-            3450521, 3450520, 3450519, 3450518, 3450517, 3450514, 3450513,
-            3450512, 3450511, 3450510, 3450509, 3450508, 3450507, 3450506,
-            3450505, 3450503, 3450501, 3450498, 3450496, 3450493, 3450483,
-            3450481, 3450480, 3450479, 3450478, 3450477, 3450476, 3450475,
-            3450474, 3450473, 3450472, 3450471, 3450470, 3450469, 3450463,
-            3450458, 3450455, 3450451, 3450441, 3450440, 3450439, 3450438,
-            3450437, 3450436, 3450435, 3450434, 3450433, 3450432, 3450431,
-            3450430, 3450429, 3450428,
-        ],
+        "account": "l11223344",
+        "xc_token": "e1c0db2959e5411f9342c8550b03f6e9",
     },
 ]
 
 _shutdown = threading.Event()
 _state: dict[str, Any] = {
     "strategy_id": STRATEGY_ID,
+    "target_gpu": TARGET_GPU,
+    "target_status": TARGET_STATUS,
     "accounts": {
         target["account"]: {
-            "total": len(target["task_ids"]),
+            "phase": "pending",
+            "total": 0,
             "stopped": 0,
             "failed": 0,
             "failed_task_ids": [],
@@ -103,7 +61,7 @@ _state: dict[str, Any] = {
         for target in TARGETS
     },
     "phase": "starting",  # starting | stopping | done | partial_failure | error
-    "total": sum(len(target["task_ids"]) for target in TARGETS),
+    "total": 0,
     "stopped": 0,
     "failed": 0,
     "failed_task_ids": [],
@@ -147,11 +105,69 @@ def _run_http() -> None:
     server.server_close()
 
 
+def _fetch_target_task_ids(account: str, xc_token: str) -> list[int]:
+    """Return only this account's waiting tasks on the exact target GPU."""
+    headers = {"xc-Token": xc_token}
+    task_ids: list[int] = []
+    page = 1
+
+    while True:
+        response = requests.get(
+            f"{BASE_URL}{TASK_PAGE_ENDPOINT}",
+            headers=headers,
+            params={
+                "current": page,
+                "pageSize": PAGE_SIZE,
+                "status": TARGET_STATUS,
+                "gpuType": TARGET_GPU,
+            },
+            timeout=REQUEST_TIMEOUT,
+        )
+        response.raise_for_status()
+        result = response.json()
+        if result.get("code") != 0:
+            raise RuntimeError(
+                f"{account}: task query failed: {result.get('message', 'unknown error')}"
+            )
+
+        data = result.get("data") or {}
+        for record in data.get("records") or []:
+            record_account = record.get("userAccount")
+            record_status = str(record.get("status", "")).lower()
+            record_gpu = record.get("gpuType")
+            if (
+                record_account != account
+                or record_status != TARGET_STATUS
+                or record_gpu != TARGET_GPU
+            ):
+                print(
+                    f"[query] {account}: rejected mismatched record "
+                    f"taskId={record.get('taskId')} account={record_account!r} "
+                    f"status={record_status!r} gpuType={record_gpu!r}",
+                    flush=True,
+                )
+                continue
+            try:
+                task_ids.append(int(record["taskId"]))
+            except (KeyError, TypeError, ValueError):
+                print(f"[query] {account}: rejected non-numeric task ID", flush=True)
+
+        pages = int(data.get("pages") or 0)
+        if page >= pages:
+            break
+        page += 1
+
+    unique_task_ids = list(dict.fromkeys(task_ids))
+    print(
+        f"[query] {account}: found {len(unique_task_ids)} "
+        f"{TARGET_STATUS} tasks on {TARGET_GPU}",
+        flush=True,
+    )
+    return unique_task_ids
+
+
 def _stop_batch(account: str, xc_token: str, task_ids: list[int]) -> bool:
-    headers = {
-        "Content-Type": "application/json",
-        "xc-Token": xc_token,
-    }
+    headers = {"Content-Type": "application/json", "xc-Token": xc_token}
     url = f"{BASE_URL}{STOP_TASK_ENDPOINT}"
     for attempt in range(1, MAX_RETRIES + 1):
         try:
@@ -167,17 +183,20 @@ def _stop_batch(account: str, xc_token: str, task_ids: list[int]) -> bool:
                 result = {"message": response.text[:500]}
 
             if response.ok and result.get("code") == 0:
-                print(f"[stop] {account}: stopped task IDs: {task_ids}", flush=True)
+                print(
+                    f"[stop] {account}: stopped {len(task_ids)} task IDs",
+                    flush=True,
+                )
                 return True
             print(
-                f"[stop] {account}: attempt {attempt}/{MAX_RETRIES} failed for {task_ids}: "
-                f"HTTP {response.status_code}, {result}",
+                f"[stop] {account}: attempt {attempt}/{MAX_RETRIES} failed "
+                f"for {len(task_ids)} tasks: HTTP {response.status_code}, {result}",
                 flush=True,
             )
         except requests.RequestException as exc:
             print(
                 f"[stop] {account}: attempt {attempt}/{MAX_RETRIES} "
-                f"request error for {task_ids}: {exc}",
+                f"request error: {exc}",
                 flush=True,
             )
 
@@ -193,9 +212,16 @@ def _run_worker() -> None:
     _state["phase"] = "stopping"
     try:
         for target in TARGETS:
+            if _shutdown.is_set():
+                break
             account = target["account"]
             account_state = _state["accounts"][account]
-            task_ids = target["task_ids"]
+            account_state["phase"] = "querying"
+            task_ids = _fetch_target_task_ids(account, target["xc_token"])
+            account_state["total"] = len(task_ids)
+            _state["total"] += len(task_ids)
+            account_state["phase"] = "stopping"
+
             for start in range(0, len(task_ids), BATCH_SIZE):
                 if _shutdown.is_set():
                     break
@@ -209,6 +235,10 @@ def _run_worker() -> None:
                     account_state["failed"] += len(batch)
                     account_state["failed_task_ids"].extend(batch)
 
+            account_state["phase"] = (
+                "done" if account_state["failed"] == 0 else "partial_failure"
+            )
+
         _state["phase"] = "done" if _state["failed"] == 0 else "partial_failure"
     except Exception as exc:  # exposed through /status for diagnosis
         _state["phase"] = "error"