forked from luopingyi/xc_agent_platform_demo
50 lines
2.9 KiB
Python
50 lines
2.9 KiB
Python
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
from src.pending_pool import build_ordered_candidates, parse_pending_txt
|
||
|
|
|
||
|
|
|
||
|
|
class TestPendingPool(unittest.TestCase):
|
||
|
|
def test_order_filters_and_first_gpu(self):
|
||
|
|
with tempfile.TemporaryDirectory() as directory:
|
||
|
|
root = Path(directory)
|
||
|
|
md, txt, dnr, banned = [root / n for n in ('pending.md', 'pending.txt', 'dnr.md', 'banned.md')]
|
||
|
|
md.write_text('''## Section B
|
||
|
|
| 1 | `org/cold` | 50 | | |
|
||
|
|
## Section A
|
||
|
|
| 1 | `org/good` | 曦望 | 壁砺, 天垓150 | — |
|
||
|
|
| 2 | `org/exact` | 曦望 | 壁砺 | — |
|
||
|
|
| 3 | `org/Heretic-model` | 曦望 | 壁砺 | — |
|
||
|
|
| 4 | `org/OLMo-2-model` | 曦望 | 壁砺 | — |
|
||
|
|
''')
|
||
|
|
dnr.write_text('| 1 | `org/exact` | fail |\n## Uniqueness collisions\n- `org/collision`\n## Overnight\n- `org/overnight`')
|
||
|
|
banned.write_text('`abliterated` `heretic` `uncensored` `unsensored`')
|
||
|
|
rows = build_ordered_candidates(md, txt, dnr, banned)
|
||
|
|
self.assertEqual([r['section'] for r in rows], ['A'] * 4 + ['B'])
|
||
|
|
self.assertEqual(rows[0]['try_next_gpus'][0], '壁砺')
|
||
|
|
self.assertEqual(rows[1]['skip_reason'], 'do-not-resubmit')
|
||
|
|
self.assertEqual(rows[2]['status'], 'skipped')
|
||
|
|
self.assertIn('family', rows[3]['skip_reason'])
|
||
|
|
self.assertTrue(rows[-1]['needs_human'])
|
||
|
|
txt.write_text('org/b\t#B cold\norg/a\t#A done:曦望; still try:壁砺+天垓150\norg/c\\t#A try:天垓150\n')
|
||
|
|
self.assertEqual(parse_pending_txt(txt)[1]['try_next_gpus'], ['壁砺', '天垓150'])
|
||
|
|
fallback = build_ordered_candidates(root/'absent', txt, dnr, banned)
|
||
|
|
self.assertEqual([r['model_address'] for r in fallback], ['org/a', 'org/c', 'org/b'])
|
||
|
|
self.assertEqual(len(build_ordered_candidates(md, txt, dnr, banned)), 5)
|
||
|
|
|
||
|
|
def test_experience_order_missing_only_and_ties(self):
|
||
|
|
from src.pending_pool import experience_order, parse_pending_md
|
||
|
|
self.assertEqual(experience_order(['智铠 mrv-100', '天垓100', '曦望', '天垓150', '壁砺']),
|
||
|
|
['曦望', '壁砺', '天垓150', '天垓100', '智铠 mrv-100'])
|
||
|
|
self.assertEqual(experience_order(['天垓150']), ['天垓150'])
|
||
|
|
with tempfile.TemporaryDirectory() as directory:
|
||
|
|
path = Path(directory) / 'pending.md'
|
||
|
|
path.write_text('## Section A\n| 1 | `org/a` | 壁砺:验证失败; 曦望:正在验证 | 天垓150 | no 壁砺 retry |\n## Section B\n| 1 | `org/b` | 0 | 壁砺 |\n')
|
||
|
|
from unittest.mock import patch
|
||
|
|
import os
|
||
|
|
with patch.dict(os.environ, {'PENDING_SECTION': 'A'}):
|
||
|
|
rows = build_ordered_candidates(path, dnr_path=Path(directory)/'none', banned_path=Path(directory)/'none')
|
||
|
|
self.assertEqual(len(rows), 1)
|
||
|
|
self.assertEqual(rows[0]['try_next_gpus'], ['天垓150'])
|
||
|
|
self.assertEqual(rows[0]['perf_order'], 'unset')
|