xc_agent_platform_demo/src/live_cycle.py

"""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, read_queue_stats, model_key
from .pending_pool import build_ordered_candidates
from .task_add_client import TaskAddClient, build_payload, resolve_token_with_source, token_probe_status
from .ledger import load_baked
from .marks import load_marks, REJECT_CODES, PERMANENT_PAIR_CODES
from .precheck import precheck_model, SAFE_STATUSES

_submitted = set()
_marked = {}
_cycle_lock = threading.Lock()


def run_cycle(candidates=None, client=None, queue_reader=None, now=None, stats_reader=None, precheck=None):
    with _cycle_lock:
        return _run_cycle(candidates, client, queue_reader, now, stats_reader, precheck)


def _run_cycle(candidates, client, queue_reader, now, stats_reader, precheck):
    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:
        _, source = resolve_token_with_source()
    except (ValueError, OSError, UnicodeError):
        statuses = token_probe_status()
        print('[live] token_probe ' + ' '.join(
            name + '=' + statuses[name] for name in
            ('EXTERNAL_SERVICE_TOKEN', 'XC_TOKEN_FILE', 'XC_TOKEN', 'MODELHUB_XC_TOKEN')), file=sys.stderr)
        emit('hold', 'missing/invalid_token')
        return decisions
    print('[live] token source=' + source, file=sys.stderr)
    # An injected queue reader keeps existing offline callers offline.
    reason, stats = 'not_read', None
    if stats_reader is not None or queue_reader is None:
        try:
            _, reason, _, stats = (stats_reader or read_queue_stats)()
        except Exception:
            reason, stats = 'read_failed', None
    if stats is not None:
        try:
            line = ' '.join(f'{key}={stats[key]:g}' if key == 'points'
                            else f'{key}={int(stats[key])}' for key in
                            ('total', 'queue', 'verify_pass', 'verify_fail', 'failed', 'unknown', 'points'))
        except (KeyError, TypeError, ValueError, OverflowError):
            stats, reason = None, 'invalid_stats'
    if stats is None:
        # Only known reasons may enter logs; injected errors may contain secrets.
        safe_reasons = {'not_read', 'read_failed', 'invalid_stats', 'missing/invalid_token',
                        'http_error', 'malformed_json', 'transport/http_error', 'platform_error',
                        'malformed_data', 'malformed_records', 'pagination_total_changed',
                        'pagination_inconsistent', 'incomplete_pagination'}
        line = 'unavailable reason=' + (reason if isinstance(reason, str) and reason in safe_reasons else 'read_failed')
    print('[live] stats ' + line, file=sys.stderr)
    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
    try:
        attempt_cap = int(os.environ.get('LIVE_MAX_ATTEMPTS_PER_CYCLE', '15'))
        if attempt_cap < 1:
            raise ValueError
    except ValueError:
        attempt_cap = 15
    if candidates is None:
        try:
            if os.environ.get('PENDING_SUBMIT_MD') or os.environ.get('PENDING_SUBMIT_TXT'):
                candidates = build_ordered_candidates()
            else:
                candidates = json.loads((Path(__file__).parent / 'data/candidates.json').read_text())
        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
    try:
        marks = load_marks()
    except ValueError:
        emit('hold', 'marks_read_failed')
        return decisions
    try:
        baked, blocked = load_baked()
        known = baked | _submitted
    except (OSError, ValueError, KeyError, TypeError):
        emit('hold', 'state_read_failed')
        return decisions
    try:
        precheck_cap = int(os.environ.get('LIVE_MAX_PRECHECKS_PER_CYCLE', '30'))
        if precheck_cap < 1:
            raise ValueError
    except ValueError:
        precheck_cap = 30
    prechecks = 0
    emitted_marks = set()

    def choose(choices, key, model, mapping):
        selected = None
        for choice in choices:
            if choice not in GPU_ENUMS:
                continue
            enum = GPU_ENUMS[choice]
            pair = (key, enum)
            if pair in known:
                continue
            code = marks.get(pair, _marked.get(pair))
            reason = 'marked:' + str(code) if code is not None else None
            if enum in mapping:
                code = mapping[enum]
                _marked[pair] = code
                if pair not in marks:
                    reason = 'marked:precheck_' + str(code)
            if reason:
                if pair not in emitted_marks:
                    emit('skip', reason, model=model, gpu=enum)
                    emitted_marks.add(pair)
                continue
            if selected is None:
                selected = choice
        return selected

    reader = queue_reader or read_queue
    client = client if client is not None else TaskAddClient()
    attempts = 0
    counted = 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
        if not model or '<' in model or '>' in model:
            emit('skip', 'placeholder')
            continue
        key = model_key(model)
        if key in blocked:
            emit('skip', 'do-not-resubmit', model=model)
            continue
        if candidate.get('skip_reason'):
            emit('skip', 'pending_filter', model=model)
            continue
        if counted >= cap:
            emit('hold', 'cycle_cap')
            break
        mapping = {}
        choices = candidate.get('try_next_gpus') or [candidate.get('targetGpu')]
        gpu = choose(choices, key, model, mapping)
        if gpu is None:
            emit('skip', 'all_gpus_done_or_unmapped', model=model)
            continue
        if os.environ.get('LIVE_PRECHECK', '1') != '0':
            if prechecks >= precheck_cap:
                emit('hold', 'precheck_cap')
                break
            prechecks += 1
            try:
                status, mapping = (precheck or precheck_model)(model)
                if status not in SAFE_STATUSES:
                    status = 'malformed_data'
                if status == 'ok' and not isinstance(mapping, dict):
                    status = 'malformed_data'
                if status != 'ok':
                    mapping = {}
                mapping = {g: c for g, c in mapping.items()
                           if g in GPU_ENUMS.values() and type(c) is int and c in (60027, 60028)}
            except Exception:
                status, mapping = 'transport_error', {}
            if status != 'ok':
                emit('precheck', status, model=model)
            gpu = choose(choices, key, model, mapping)
            if gpu is None:
                emit('skip', 'all_gpus_done_or_unmapped', model=model)
                continue
        # Rebuild remote identities before config calls and every submission.
        try:
            count, why, remote = reader()
            if type(count) is not int or count < 0:
                raise ValueError
            known.update(remote)
        except Exception:
            emit('hold', 'ledger_rebuild_failed')
            break
        gpu = choose(choices, key, model, mapping)
        if gpu is None:
            emit('skip', 'all_gpus_done_or_unmapped', model=model)
            continue
        if count >= queue_threshold():
            emit('hold', f'queue>={queue_threshold()}', count=count)
            break
        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
        pair = (key, GPU_ENUMS[gpu])
        known.add(pair)
        _submitted.add(pair)
        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'
        if code is not None and code != 0:
            emit('rejected', REJECT_CODES.get(code, 'unknown'), model=model,
                 gpu=GPU_ENUMS[gpu], code=code, http=status, queue_before=count)
            _submitted.discard(pair)
            known.discard(pair)
            if code in PERMANENT_PAIR_CODES or code != 60007:
                _marked[pair] = code
            if code == 60007:
                emit('hold', 'task_num_limit')
                break
            if attempts >= attempt_cap:
                emit('hold', 'attempt_cap')
                break
            continue
        counted += 1
        emit('submitted', reason, model=model, gpu=GPU_ENUMS[gpu], queue_before=count,
             http=status, code=code, **meta)
        if attempts >= attempt_cap:
            emit('hold', 'attempt_cap')
            break
    return decisions