forked from luopingyi/xc_agent_platform_demo
138 lines
5.9 KiB
Python
138 lines
5.9 KiB
Python
|
|
"""Queue-gated live cycles; all remote collaborators are injectable."""
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import threading
|
||
|
|
from pathlib import Path
|
||
|
|
from .banned import banned_rule
|
||
|
|
from .config_params import GPU_ENUMS, build_config_params_with_meta
|
||
|
|
from .gate import evaluate_gate, queue_threshold
|
||
|
|
from .live_queue import read_queue, model_key
|
||
|
|
from .pending_pool import build_ordered_candidates
|
||
|
|
from .task_add_client import TaskAddClient, build_payload
|
||
|
|
|
||
|
|
_submitted = set()
|
||
|
|
_cycle_lock = threading.Lock()
|
||
|
|
|
||
|
|
|
||
|
|
def run_cycle(candidates=None, client=None, queue_reader=None, now=None):
|
||
|
|
with _cycle_lock:
|
||
|
|
return _run_cycle(candidates, client, queue_reader, now)
|
||
|
|
|
||
|
|
|
||
|
|
def _run_cycle(candidates, client, queue_reader, now):
|
||
|
|
decisions = []
|
||
|
|
|
||
|
|
def emit(action, reason, **fields):
|
||
|
|
decision = dict(action=action, reason=reason, **fields)
|
||
|
|
decisions.append(decision)
|
||
|
|
print('[live] ' + action + ' ' + ' '.join(f'{k}={str(v).replace(chr(10), " ").replace(chr(13), " ")}' for k, v in dict(fields, reason=reason).items()), file=sys.stderr)
|
||
|
|
|
||
|
|
gate = evaluate_gate()
|
||
|
|
if not gate.allowed:
|
||
|
|
emit('hold', gate.reason)
|
||
|
|
return decisions
|
||
|
|
try:
|
||
|
|
cap = int(os.environ.get('LIVE_MAX_SUBMITS_PER_CYCLE', '3'))
|
||
|
|
if cap < 1:
|
||
|
|
raise ValueError
|
||
|
|
except ValueError:
|
||
|
|
emit('hold', 'invalid_cycle_cap')
|
||
|
|
return decisions
|
||
|
|
if candidates is None:
|
||
|
|
fallback = Path(__file__).parent / 'data/pending-submit.md'
|
||
|
|
try:
|
||
|
|
candidates = build_ordered_candidates(
|
||
|
|
os.environ.get('PENDING_SUBMIT_MD', str(fallback)),
|
||
|
|
os.environ.get('PENDING_SUBMIT_TXT', str(fallback.with_suffix('.txt'))))
|
||
|
|
except (OSError, ValueError, IndexError):
|
||
|
|
emit('hold', 'candidate_read_failed')
|
||
|
|
return decisions
|
||
|
|
candidates = sorted(candidates, key=lambda c: (c.get('section', 'B') != 'A', c.get('foothold_class') != 'verified_reuse'))
|
||
|
|
if not candidates:
|
||
|
|
emit('hold', 'no_candidates')
|
||
|
|
return decisions
|
||
|
|
state = Path(os.environ.get('LIVE_STATE_FILE', '/tmp/live-submitted.json'))
|
||
|
|
try:
|
||
|
|
saved = json.loads(state.read_text()) if state.exists() else []
|
||
|
|
if not isinstance(saved, list) or any(not isinstance(s, str) for s in saved):
|
||
|
|
raise ValueError
|
||
|
|
known = set(map(model_key, saved)) | _submitted
|
||
|
|
except (OSError, ValueError):
|
||
|
|
emit('hold', 'state_read_failed')
|
||
|
|
return decisions
|
||
|
|
reader = queue_reader or read_queue
|
||
|
|
client = client if client is not None else TaskAddClient()
|
||
|
|
attempts = 0
|
||
|
|
for candidate in candidates:
|
||
|
|
model = candidate.get('model_address') or candidate.get('modelAddress') or ''
|
||
|
|
rule = banned_rule(model + ' ' + str(candidate.get('model_name', candidate.get('modelName', ''))))
|
||
|
|
if rule:
|
||
|
|
emit('skip', 'banned:' + rule, model=model)
|
||
|
|
continue
|
||
|
|
key = model_key(model)
|
||
|
|
if key in known:
|
||
|
|
emit('skip', 'duplicate', model=model)
|
||
|
|
continue
|
||
|
|
if candidate.get('skip_reason'):
|
||
|
|
emit('skip', 'pending_filter', model=model)
|
||
|
|
continue
|
||
|
|
if attempts >= cap:
|
||
|
|
emit('hold', 'cycle_cap')
|
||
|
|
break
|
||
|
|
gpu = next(iter(candidate.get('try_next_gpus') or []), candidate.get('targetGpu'))
|
||
|
|
if gpu not in GPU_ENUMS:
|
||
|
|
emit('skip', 'gpu_unmapped', model=model)
|
||
|
|
continue
|
||
|
|
framework = candidate.get('framework') or 'vllm'
|
||
|
|
try:
|
||
|
|
yaml, meta = build_config_params_with_meta(gpu, framework)
|
||
|
|
except Exception:
|
||
|
|
emit('skip', 'config_failed', model=model)
|
||
|
|
continue
|
||
|
|
payload = build_payload(model, GPU_ENUMS[gpu], framework, os.environ['STRATEGY_ID'].strip(), yaml)
|
||
|
|
if any(not value or '<' in str(value) or '>' in str(value) for value in payload.values()):
|
||
|
|
emit('skip', 'placeholder', model=model)
|
||
|
|
continue
|
||
|
|
# Fresh complete snapshot immediately before reserving and sending each add.
|
||
|
|
try:
|
||
|
|
count, why, remote = reader()
|
||
|
|
except Exception:
|
||
|
|
count, why, remote = None, 'reader_error', set()
|
||
|
|
if count is None:
|
||
|
|
# Only internal fixed reasons are emitted; injected errors can contain secrets.
|
||
|
|
safe = why if why in {'missing/invalid_token', 'http_error', 'malformed_json', 'transport/http_error', 'platform_error', 'malformed_data', 'malformed_records', 'pagination_total_changed', 'pagination_inconsistent', 'incomplete_pagination'} else 'reader_error'
|
||
|
|
emit('hold', 'queue_read_failed:' + safe)
|
||
|
|
break
|
||
|
|
if type(count) is not int or count < 0:
|
||
|
|
emit('hold', 'queue_read_failed:invalid_count')
|
||
|
|
break
|
||
|
|
if count >= queue_threshold():
|
||
|
|
emit('hold', f'queue>={queue_threshold()}', count=count)
|
||
|
|
break
|
||
|
|
if key in {model_key(m) for m in remote}:
|
||
|
|
emit('skip', 'duplicate', model=model)
|
||
|
|
continue
|
||
|
|
known.add(key)
|
||
|
|
_submitted.add(key)
|
||
|
|
try:
|
||
|
|
state.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary = state.with_name(state.name + '.tmp')
|
||
|
|
temporary.write_text(json.dumps(sorted(known)))
|
||
|
|
temporary.replace(state)
|
||
|
|
except OSError:
|
||
|
|
emit('hold', 'state_write_failed')
|
||
|
|
break
|
||
|
|
attempts += 1
|
||
|
|
try:
|
||
|
|
result = client.submit(payload)
|
||
|
|
status, code = result.get('http_status'), result.get('platform_code')
|
||
|
|
status = status if type(status) is int else None
|
||
|
|
code = code if type(code) is int else None
|
||
|
|
reason = 'task_add_response' if status is not None else 'transport_error'
|
||
|
|
except Exception:
|
||
|
|
status, code, reason = None, None, 'submit_error'
|
||
|
|
emit('submitted', reason, model=model, gpu=GPU_ENUMS[gpu], queue_before=count,
|
||
|
|
http=status, code=code, **meta)
|
||
|
|
return decisions
|