271 lines
8.8 KiB
Python
271 lines
8.8 KiB
Python
"""Stop waiting Sunrise_pt-200-x1 validation tasks for selected accounts.
|
|
|
|
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
|
|
from datetime import datetime, timezone
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
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"))
|
|
|
|
TARGETS = [
|
|
{
|
|
"account": "jiangxiaowen",
|
|
"xc_token": "88d5fee9f1fe4f7583f11a9d3702dc85",
|
|
},
|
|
{
|
|
"account": "l112233",
|
|
"xc_token": "40cb6910dc9a442a816298a228da65ac",
|
|
},
|
|
{
|
|
"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"]: {
|
|
"phase": "pending",
|
|
"total": 0,
|
|
"stopped": 0,
|
|
"failed": 0,
|
|
"failed_task_ids": [],
|
|
}
|
|
for target in TARGETS
|
|
},
|
|
"phase": "starting", # starting | stopping | done | partial_failure | error
|
|
"total": 0,
|
|
"stopped": 0,
|
|
"failed": 0,
|
|
"failed_task_ids": [],
|
|
"started_at": None,
|
|
"finished_at": None,
|
|
"error": None,
|
|
}
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def do_GET(self) -> None: # noqa: N802 - BaseHTTPRequestHandler API
|
|
if self.path == "/health":
|
|
self._json({"status": "ok"})
|
|
elif self.path == "/status":
|
|
self._json(_state)
|
|
else:
|
|
self._json({"error": "not found"}, 404)
|
|
|
|
def _json(self, body: dict[str, Any], status: int = 200) -> None:
|
|
payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
|
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: str, *args: object) -> None:
|
|
print(f"[http] {self.address_string()} {fmt % args}", flush=True)
|
|
|
|
|
|
def _run_http() -> None:
|
|
server = ThreadingHTTPServer((HTTP_HOST, HTTP_PORT), Handler)
|
|
server.timeout = 1
|
|
print(f"[http] listening on {HTTP_HOST}:{HTTP_PORT}", flush=True)
|
|
while not _shutdown.is_set():
|
|
server.handle_request()
|
|
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}
|
|
url = f"{BASE_URL}{STOP_TASK_ENDPOINT}"
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
response = requests.put(
|
|
url,
|
|
headers=headers,
|
|
json={"taskIds": task_ids},
|
|
timeout=REQUEST_TIMEOUT,
|
|
)
|
|
try:
|
|
result = response.json()
|
|
except ValueError:
|
|
result = {"message": response.text[:500]}
|
|
|
|
if response.ok and result.get("code") == 0:
|
|
print(
|
|
f"[stop] {account}: stopped {len(task_ids)} task IDs",
|
|
flush=True,
|
|
)
|
|
return True
|
|
print(
|
|
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: {exc}",
|
|
flush=True,
|
|
)
|
|
|
|
if attempt < MAX_RETRIES and not _shutdown.wait(attempt):
|
|
continue
|
|
if _shutdown.is_set():
|
|
break
|
|
return False
|
|
|
|
|
|
def _run_worker() -> None:
|
|
_state["started_at"] = _now()
|
|
_state["phase"] = "stopping"
|
|
try:
|
|
for target in TARGETS:
|
|
if _shutdown.is_set():
|
|
break
|
|
account = target["account"]
|
|
account_state = _state["accounts"][account]
|
|
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
|
|
batch = task_ids[start : start + BATCH_SIZE]
|
|
if _stop_batch(account, target["xc_token"], batch):
|
|
_state["stopped"] += len(batch)
|
|
account_state["stopped"] += len(batch)
|
|
else:
|
|
_state["failed"] += len(batch)
|
|
_state["failed_task_ids"].extend(batch)
|
|
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"
|
|
_state["error"] = str(exc)
|
|
print(f"[stop] fatal error: {exc}", flush=True)
|
|
finally:
|
|
_state["finished_at"] = _now()
|
|
print(f"[stop] completed: {_state}", flush=True)
|
|
|
|
|
|
def _handle_signal(signum: int, _frame: Any) -> None:
|
|
print(f"[main] received signal {signum}; shutting down", flush=True)
|
|
_shutdown.set()
|
|
|
|
|
|
def main() -> None:
|
|
signal.signal(signal.SIGTERM, _handle_signal)
|
|
signal.signal(signal.SIGINT, _handle_signal)
|
|
|
|
http_thread = threading.Thread(target=_run_http, daemon=False)
|
|
http_thread.start()
|
|
threading.Thread(target=_run_worker, daemon=True).start()
|
|
|
|
_shutdown.wait()
|
|
http_thread.join(timeout=5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|