forked from luopingyi/xc_agent_platform_demo
225 lines
12 KiB
Python
225 lines
12 KiB
Python
import io
|
|
import json
|
|
from pathlib import Path
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
from src.live_cycle import run_cycle, _submitted, _rejected
|
|
|
|
|
|
def candidate(name, section='B', foothold='unknown'):
|
|
return dict(model_address='org/' + name, section=section, foothold_class=foothold, try_next_gpus=['壁砺'])
|
|
|
|
|
|
class TestLiveCycle(unittest.TestCase):
|
|
def setUp(self):
|
|
directory = tempfile.TemporaryDirectory()
|
|
self.addCleanup(directory.cleanup)
|
|
self.env = patch.dict(os.environ, {'STRATEGY_ID': 'sid', 'LIVE_STATE_FILE': directory.name + '/state.json', 'CONFIG_PARAMS_SOURCE': 'sample', 'XC_TOKEN': 'synthetic-secret'}, clear=True)
|
|
self.env.start()
|
|
self.addCleanup(self.env.stop)
|
|
_submitted.clear()
|
|
_rejected.clear()
|
|
self.client = Mock()
|
|
self.client.submit.return_value = dict(http_status=200, platform_code=0)
|
|
network = patch('urllib.request.urlopen', side_effect=AssertionError('network forbidden'))
|
|
network.start()
|
|
self.addCleanup(network.stop)
|
|
logs = patch('sys.stderr', new_callable=io.StringIO)
|
|
self.logs = logs.start()
|
|
self.addCleanup(logs.stop)
|
|
|
|
def cycle(self, rows=None, counts=(94,)):
|
|
reader = Mock(side_effect=[(n, 'platform_error' if n is None else 'ok', set()) for n in counts])
|
|
return run_cycle(rows or [candidate('a')], self.client, reader)
|
|
|
|
def state_pairs(self):
|
|
return json.loads(Path(os.environ['LIVE_STATE_FILE']).read_text())
|
|
|
|
def test_done_rejection_keeps_reservation_and_continues(self):
|
|
for code in (60027, 60028):
|
|
with self.subTest(code=code):
|
|
self.client.submit.side_effect = [
|
|
dict(http_status=200, platform_code=code),
|
|
dict(http_status=200, platform_code=0)]
|
|
with patch.dict(os.environ, {'LIVE_MAX_SUBMITS_PER_CYCLE': '1'}):
|
|
decisions = self.cycle([candidate(str(code)), candidate('ok' + str(code))], (1, 1))
|
|
self.assertEqual([d['action'] for d in decisions], ['rejected', 'submitted'])
|
|
self.assertIn('org/' + str(code) + '|Biren_166m', self.state_pairs())
|
|
self.assertIn('rejected', self.logs.getvalue())
|
|
self.assertIn('CREATE_CONTEST_TASK_MODEL_VERIFY_RUNNING', self.logs.getvalue())
|
|
|
|
def test_rejection_removed_but_not_retried_in_process(self):
|
|
self.client.submit.return_value = dict(http_status=200, platform_code=60020)
|
|
self.cycle()
|
|
self.assertEqual(self.state_pairs(), [])
|
|
self.cycle(counts=())
|
|
self.client.submit.assert_called_once()
|
|
|
|
def test_unknown_rejection(self):
|
|
self.client.submit.return_value = dict(http_status=200, platform_code=99999)
|
|
decisions = self.cycle()
|
|
self.assertEqual(decisions[0]['action'], 'rejected')
|
|
self.assertEqual(decisions[0]['reason'], 'unknown')
|
|
self.assertEqual(self.state_pairs(), [])
|
|
|
|
def test_attempt_cap_and_invalid_fallback(self):
|
|
for value, expected in ((None, 15), ('bad', 15), ('0', 15), ('-1', 15), ('2', 2)):
|
|
with self.subTest(value=value):
|
|
_rejected.clear()
|
|
self.client.reset_mock()
|
|
self.client.submit.return_value = dict(http_status=200, platform_code=60020)
|
|
with patch.dict(os.environ, {} if value is None else {'LIVE_MAX_ATTEMPTS_PER_CYCLE': value}):
|
|
decisions = self.cycle([candidate(str(i)) for i in range(expected)], (1,) * expected)
|
|
self.assertEqual(self.client.submit.call_count, expected)
|
|
self.assertEqual(decisions[-1]['reason'], 'attempt_cap')
|
|
|
|
def test_submit_cap_counts_successes_after_rejections(self):
|
|
self.client.submit.side_effect = [dict(http_status=200, platform_code=c) for c in (60020, 0, 60028, 0, 0)]
|
|
decisions = self.cycle([candidate(str(i)) for i in range(7)], (1,) * 5)
|
|
self.assertEqual(self.client.submit.call_count, 5)
|
|
self.assertEqual(sum(d['action'] == 'submitted' for d in decisions), 3)
|
|
self.assertEqual(decisions[-1]['reason'], 'cycle_cap')
|
|
|
|
def test_task_limit_stops_cycle(self):
|
|
self.client.submit.return_value = dict(http_status=200, platform_code=60007)
|
|
decisions = self.cycle([candidate('a'), candidate('b')])
|
|
self.assertEqual([d['action'] for d in decisions], ['rejected', 'hold'])
|
|
self.assertEqual(decisions[-1]['reason'], 'task_num_limit')
|
|
self.assertEqual(self.state_pairs(), [])
|
|
self.client.submit.assert_called_once()
|
|
|
|
def test_unknown_outcomes_count_and_reserve_before_submit(self):
|
|
for result in (dict(http_status=200, platform_code='0'), dict(http_status=None, platform_code=None), OSError()):
|
|
with self.subTest(result=result):
|
|
_submitted.clear()
|
|
Path(os.environ['LIVE_STATE_FILE']).unlink(missing_ok=True)
|
|
self.client.reset_mock()
|
|
def submit(payload):
|
|
self.assertIn('org/a|Biren_166m', self.state_pairs())
|
|
if isinstance(result, Exception):
|
|
raise result
|
|
return result
|
|
self.client.submit.side_effect = submit
|
|
with patch.dict(os.environ, {'LIVE_MAX_SUBMITS_PER_CYCLE': '1'}):
|
|
decisions = self.cycle([candidate('a'), candidate('b')])
|
|
self.assertEqual(decisions[0]['action'], 'submitted')
|
|
self.assertEqual(decisions[-1]['reason'], 'cycle_cap')
|
|
self.client.submit.assert_called_once()
|
|
|
|
def test_high_and_failed_queue(self):
|
|
for n in (95, None):
|
|
decisions = self.cycle(counts=(n,))
|
|
self.assertEqual(decisions[-1]['action'], 'hold')
|
|
self.client.submit.assert_not_called()
|
|
self.assertIn('[live] hold', self.logs.getvalue())
|
|
|
|
def test_order_and_cap(self):
|
|
rows = [candidate('b'), candidate('a', 'A'), candidate('verified', 'A', 'verified_reuse'), candidate('c'), candidate('d')]
|
|
self.cycle(rows, (90, 90, 90))
|
|
self.assertEqual([c.args[0]['modelAddress'] for c in self.client.submit.call_args_list], ['org/verified', 'org/a', 'org/b'])
|
|
self.assertIn('cycle_cap', self.logs.getvalue())
|
|
|
|
def test_midcycle(self):
|
|
self.cycle([candidate(str(i)) for i in range(5)], (93, 94, 95))
|
|
self.assertEqual(self.client.submit.call_count, 2)
|
|
|
|
def test_kill_and_threshold(self):
|
|
for env in ({'ALLOW_LIVE_SUBMIT': '0'}, {'LIVE_QUEUE_THRESHOLD': 'bad'}, {'LIVE_QUEUE_THRESHOLD': '20'}):
|
|
with patch.dict(os.environ, env):
|
|
self.cycle(counts=(20,))
|
|
self.client.submit.assert_not_called()
|
|
with patch.dict(os.environ, {'LIVE_QUEUE_THRESHOLD': '100'}):
|
|
self.cycle(counts=(99,))
|
|
self.client.submit.assert_called_once()
|
|
|
|
def test_banned_duplicate_persistence(self):
|
|
rows = [candidate('GPT-2'), candidate('ok'), candidate('ok')]
|
|
self.cycle(rows)
|
|
_submitted.clear()
|
|
self.cycle([candidate('ok')], ())
|
|
self.client.submit.assert_called_once()
|
|
self.assertIn('banned:', self.logs.getvalue())
|
|
self.assertIn('all_gpus_done_or_unmapped', self.logs.getvalue())
|
|
|
|
def test_remote_duplicate_and_transport_reservation(self):
|
|
run_cycle([candidate('remote')], self.client, lambda: (1, 'ok', {('org/remote', 'Biren_166m')}))
|
|
self.client.submit.assert_not_called()
|
|
self.client.submit.side_effect = OSError('do not log')
|
|
self.cycle()
|
|
self.cycle(counts=())
|
|
self.client.submit.assert_called_once()
|
|
self.assertNotIn('do not log', self.logs.getvalue())
|
|
|
|
def test_missing_token_real_reader(self):
|
|
with patch.dict(os.environ, {'STRATEGY_ID': 'sid'}, clear=True):
|
|
run_cycle([candidate('a')], self.client)
|
|
self.client.submit.assert_not_called()
|
|
self.assertIn('missing/invalid_token', self.logs.getvalue())
|
|
self.assertIn('[live] token_probe EXTERNAL_SERVICE_TOKEN=absent XC_TOKEN_FILE=absent XC_TOKEN=absent MODELHUB_XC_TOKEN=absent', self.logs.getvalue())
|
|
self.assertNotIn('[live] token source=', self.logs.getvalue())
|
|
|
|
def test_rejected_token_probe_does_not_leak(self):
|
|
with patch.dict(os.environ, {'XC_TOKEN': '<private-secret-fragment>'}):
|
|
self.cycle()
|
|
self.client.submit.assert_not_called()
|
|
self.assertIn('XC_TOKEN=rejected', self.logs.getvalue())
|
|
self.assertNotIn('private-secret', self.logs.getvalue())
|
|
|
|
def test_worker_waits_and_stops(self):
|
|
from src.server import _live_worker
|
|
with patch('src.server._shutdown') as shutdown, patch('src.server.run_cycle') as cycle:
|
|
shutdown.is_set.return_value = False
|
|
shutdown.wait.return_value = True
|
|
with patch.dict(os.environ, {'LIVE_CYCLE_SECONDS': '1'}):
|
|
_live_worker()
|
|
cycle.assert_called_once()
|
|
shutdown.wait.assert_called_once_with(60)
|
|
|
|
def test_unauthenticated_platform_zero_submits(self):
|
|
import json
|
|
response = Mock()
|
|
response.__enter__ = Mock(return_value=response)
|
|
response.__exit__ = Mock(return_value=False)
|
|
response.status = 200
|
|
response.read.return_value = json.dumps({'code': 40100}).encode()
|
|
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}), patch('urllib.request.urlopen', return_value=response):
|
|
run_cycle([candidate('a')], self.client)
|
|
self.client.submit.assert_not_called()
|
|
|
|
def test_unmapped_and_placeholder(self):
|
|
rows = [dict(model_address='org/unmapped', targetGpu='unknown'), candidate('<MODEL>')]
|
|
self.cycle(rows, ())
|
|
self.client.submit.assert_not_called()
|
|
self.assertIn('all_gpus_done_or_unmapped', self.logs.getvalue())
|
|
self.assertIn('placeholder', self.logs.getvalue())
|
|
|
|
def test_stats_logged_once_and_submission_continues(self):
|
|
from src.live_queue import compute_stats
|
|
stats = compute_stats([dict(status='success', verifyResult=1, stageAndNoviceTaskCredit='5.5')])
|
|
stats_reader = Mock(return_value=(0, 'ok', set(), stats))
|
|
run_cycle([candidate('stats')], self.client, lambda: (0, 'ok', set()), stats_reader=stats_reader)
|
|
stats_reader.assert_called_once_with()
|
|
lines = [s for s in self.logs.getvalue().splitlines() if s.startswith('[live] stats ')]
|
|
self.assertEqual(lines, ['[live] stats total=1 queue=0 verify_pass=1 verify_fail=0 failed=0 unknown=0 points=5.5'])
|
|
self.assertNotIn('synthetic-secret', self.logs.getvalue())
|
|
self.client.submit.assert_called_once()
|
|
|
|
def test_stats_failure_is_safe_and_does_not_gate(self):
|
|
for index, reader in enumerate((Mock(return_value=(None, 'platform_error', set(), None)),
|
|
Mock(side_effect=OSError('synthetic-secret')),
|
|
Mock(return_value=(None, 'synthetic-secret', set(), None)))):
|
|
self.logs.truncate(0)
|
|
self.logs.seek(0)
|
|
run_cycle([candidate('stats-failure-' + str(index))], self.client, lambda: (0, 'ok', set()), stats_reader=reader)
|
|
lines = [s for s in self.logs.getvalue().splitlines() if s.startswith('[live] stats ')]
|
|
self.assertEqual(len(lines), 1)
|
|
self.assertTrue(lines[0].startswith('[live] stats unavailable reason='))
|
|
self.assertNotIn('synthetic-secret', self.logs.getvalue())
|
|
self.assertEqual(self.client.submit.call_count, 3)
|
|
|
|
def test_injected_queue_skips_stats_network(self):
|
|
self.cycle()
|
|
self.assertEqual(self.logs.getvalue().count('[live] stats unavailable reason=not_read'), 1)
|