forked from luopingyi/xc_agent_platform_demo
72 lines
3.5 KiB
Python
72 lines
3.5 KiB
Python
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
from src.live_cycle import run_cycle, _submitted
|
|
from src.ledger import load_baked, rebuild
|
|
from src.banned import banned_rule
|
|
from src.task_add_client import resolve_token_with_source
|
|
|
|
|
|
class TestLedger(unittest.TestCase):
|
|
def setUp(self):
|
|
temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temp.cleanup)
|
|
self.state = Path(temp.name) / 'state.json'
|
|
env = patch.dict(os.environ, dict(STRATEGY_ID='sid', XC_TOKEN='synthetic-credential', CONFIG_PARAMS_SOURCE='sample', LIVE_STATE_FILE=str(self.state)), clear=True)
|
|
env.start()
|
|
self.addCleanup(env.stop)
|
|
_submitted.clear()
|
|
self.addCleanup(_submitted.clear)
|
|
net = patch('urllib.request.urlopen', side_effect=AssertionError('network forbidden'))
|
|
net.start()
|
|
self.addCleanup(net.stop)
|
|
self.client = Mock()
|
|
self.client.submit.return_value = dict(http_status=200, platform_code=0)
|
|
|
|
def cycle(self, remote=(), model='unit/model', gpus=None):
|
|
return run_cycle([dict(model_address=model, try_next_gpus=gpus or ['壁砺', '天垓100'])], self.client, lambda: (0, 'ok', set(remote)))
|
|
|
|
def test_pair_and_restart(self):
|
|
self.cycle([('unit/model', 'Biren_166m')])
|
|
self.assertEqual(self.client.submit.call_args.args[0]['targetGpu'], 'Iluvatar_bi-100')
|
|
self.assertEqual(json.loads(self.state.read_text()), ['unit/model|Iluvatar_bi-100'])
|
|
_submitted.clear()
|
|
self.cycle([('unit/model', 'Biren_166m')])
|
|
self.client.submit.assert_called_once()
|
|
|
|
def test_baked_and_blocked(self):
|
|
pairs, blocked = load_baked()
|
|
model, gpu = next(iter(pairs))
|
|
self.cycle(model=model, gpus=[gpu])
|
|
self.cycle(model=next(iter(blocked)))
|
|
self.client.submit.assert_not_called()
|
|
|
|
def test_failure_and_missing_token(self):
|
|
result = run_cycle([dict(model_address='unit/model', try_next_gpus=['壁砺'])], self.client, lambda: (None, 'failed', set()))
|
|
self.assertEqual(result[-1]['reason'], 'ledger_rebuild_failed')
|
|
with patch.dict(os.environ, {'XC_TOKEN': ''}):
|
|
self.cycle()
|
|
self.client.submit.assert_not_called()
|
|
|
|
def test_source_and_no_value_logs(self):
|
|
import io
|
|
with patch.dict(os.environ, {'EXTERNAL_SERVICE_TOKEN': 'platform-credential', 'XC_TOKEN_FILE': '/nonexistent'}), patch('sys.stderr', new_callable=io.StringIO) as log:
|
|
self.assertEqual(resolve_token_with_source(), ('platform-credential', 'EXTERNAL_SERVICE_TOKEN'))
|
|
self.cycle()
|
|
self.assertEqual(log.getvalue().count('token source=EXTERNAL_SERVICE_TOKEN'), 1)
|
|
self.assertNotIn('platform-credential', log.getvalue())
|
|
with patch.dict(os.environ, {'EXTERNAL_SERVICE_TOKEN': '<TOKEN>', 'XC_TOKEN': 'placeholder', 'MODELHUB_XC_TOKEN': 'alias-credential'}):
|
|
self.assertEqual(resolve_token_with_source()[1], 'MODELHUB_XC_TOKEN')
|
|
|
|
def test_candidates_and_rebuild(self):
|
|
rows = json.loads(Path('src/data/candidates.json').read_text())
|
|
self.assertGreater(len(rows), 0)
|
|
self.assertEqual([r['section'] for r in rows], sorted(r['section'] for r in rows))
|
|
self.assertFalse(any(banned_rule(r['model_address']) for r in rows))
|
|
count, pairs, blocked, sizes = rebuild(self.state, lambda: (0, 'ok', {('unit/model', 'Biren_166m')}))
|
|
self.assertIn(('unit/model', 'Biren_166m'), pairs)
|
|
self.assertEqual(sizes[0], 1)
|