forked from luopingyi/xc_agent_platform_demo
live gate: default-on queue-gated task/add (<95), fail-closed, cap 3/cycle, banned+dedupe, kill switch (54 tests)
This commit is contained in:
Victor Cao
@@ -1,159 +1,44 @@
|
||||
"""Gate defaults to dry-run / blocked; live only with dual confirm + STRATEGY_ID."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
"""Default live approval, kill switch, strategy and conservative hint gates."""
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from src.gate import LIVE_CONFIRM_VALUE, attempt_live_submit, evaluate_gate
|
||||
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = []
|
||||
|
||||
def submit(self, payload):
|
||||
self.calls.append(payload)
|
||||
return {"task_id": "fake-1"}
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from src.gate import evaluate_gate, attempt_live_submit, LIVE_CONFIRM_VALUE
|
||||
|
||||
|
||||
class TestGate(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.counter = mock.patch('src.gate._submit_count', 0)
|
||||
self.counter.start()
|
||||
self.addCleanup(self.counter.stop)
|
||||
self.env = mock.patch.dict(os.environ, {}, clear=True)
|
||||
self.env.start()
|
||||
self.addCleanup(self.env.stop)
|
||||
env = patch.dict(os.environ, {'STRATEGY_ID': 'sid'}, clear=True)
|
||||
env.start()
|
||||
self.addCleanup(env.stop)
|
||||
|
||||
def test_queue_and_process_limit(self):
|
||||
client = FakeClient()
|
||||
with mock.patch.dict(os.environ, {
|
||||
'STRATEGY_ID': 'sid', 'ALLOW_LIVE_SUBMIT': '1',
|
||||
'LIVE_CONFIRM': LIVE_CONFIRM_VALUE, 'ASYNC_QUEUE_DEPTH': '95'
|
||||
}):
|
||||
self.assertIn('≥95', attempt_live_submit({}, client)['reason'])
|
||||
self.assertEqual(client.calls, [])
|
||||
os.environ['ASYNC_QUEUE_DEPTH'] = '94'
|
||||
self.assertTrue(attempt_live_submit({}, client)['submitted'])
|
||||
self.assertFalse(attempt_live_submit({}, client)['submitted'])
|
||||
self.assertEqual(len(client.calls), 1)
|
||||
def test_default_live_and_optional_confirmation(self):
|
||||
self.assertTrue(evaluate_gate().allowed)
|
||||
for value in ('', 'YES', LIVE_CONFIRM_VALUE):
|
||||
with patch.dict(os.environ, {'LIVE_CONFIRM': value}):
|
||||
self.assertTrue(evaluate_gate().allowed)
|
||||
|
||||
def _clear_gate_env(self):
|
||||
return mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"STRATEGY_ID": "sid-1",
|
||||
"ALLOW_LIVE_SUBMIT": "0",
|
||||
"LIVE_CONFIRM": "",
|
||||
},
|
||||
clear=False,
|
||||
)
|
||||
def test_kill_switch_and_missing_strategy(self):
|
||||
for env in ({'ALLOW_LIVE_SUBMIT': '0'}, {'ALLOW_LIVE_SUBMIT': 'true'}, {'STRATEGY_ID': ' '}):
|
||||
with patch.dict(os.environ, env):
|
||||
self.assertFalse(evaluate_gate().allowed)
|
||||
|
||||
def test_default_dry_run(self) -> None:
|
||||
with self._clear_gate_env():
|
||||
d = evaluate_gate()
|
||||
self.assertFalse(d.allowed)
|
||||
self.assertEqual(d.mode, "dry-run")
|
||||
def test_hints_and_threshold(self):
|
||||
for value in ('95', 'invalid', '-1'):
|
||||
with patch.dict(os.environ, {'ASYNC_QUEUE_DEPTH': value}):
|
||||
self.assertFalse(evaluate_gate().allowed)
|
||||
with patch.dict(os.environ, {'ASYNC_QUEUE_DEPTH': '94'}):
|
||||
self.assertTrue(evaluate_gate().allowed)
|
||||
with patch.dict(os.environ, {'LIVE_QUEUE_THRESHOLD': 'bad'}):
|
||||
self.assertFalse(evaluate_gate().allowed)
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = Path(d) / 'hint.json'
|
||||
p.write_text('{"depth":95}')
|
||||
with patch.dict(os.environ, {'QUEUE_SNAPSHOT_JSON': str(p), 'ASYNC_QUEUE_DEPTH': '1'}):
|
||||
self.assertFalse(evaluate_gate().allowed)
|
||||
|
||||
def test_missing_strategy_blocked(self) -> None:
|
||||
env = {
|
||||
"ALLOW_LIVE_SUBMIT": "1",
|
||||
"LIVE_CONFIRM": LIVE_CONFIRM_VALUE,
|
||||
}
|
||||
# Remove STRATEGY_ID
|
||||
base = {k: v for k, v in os.environ.items() if k != "STRATEGY_ID"}
|
||||
base.update(env)
|
||||
with mock.patch.dict(os.environ, base, clear=True):
|
||||
d = evaluate_gate()
|
||||
self.assertFalse(d.allowed)
|
||||
self.assertEqual(d.mode, "blocked")
|
||||
|
||||
def test_allow_without_confirm_blocked(self) -> None:
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"STRATEGY_ID": "sid-1",
|
||||
"ALLOW_LIVE_SUBMIT": "1",
|
||||
"LIVE_CONFIRM": "YES",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
d = evaluate_gate()
|
||||
self.assertFalse(d.allowed)
|
||||
self.assertEqual(d.mode, "blocked")
|
||||
|
||||
def test_dual_gate_opens(self) -> None:
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"STRATEGY_ID": "sid-1",
|
||||
"ALLOW_LIVE_SUBMIT": "1",
|
||||
"LIVE_CONFIRM": LIVE_CONFIRM_VALUE,
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
d = evaluate_gate()
|
||||
self.assertTrue(d.allowed)
|
||||
self.assertEqual(d.mode, "live")
|
||||
|
||||
def test_open_gate_without_client_still_blocked(self) -> None:
|
||||
"""Phase2: an open gate never constructs a network client."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"STRATEGY_ID": "sid-1",
|
||||
"ALLOW_LIVE_SUBMIT": "1",
|
||||
"LIVE_CONFIRM": LIVE_CONFIRM_VALUE,
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
result = attempt_live_submit({"modelAddress": "x/y"})
|
||||
self.assertFalse(result["submitted"])
|
||||
self.assertIn("no real", result["reason"].lower())
|
||||
self.assertEqual(result["mode"], "blocked")
|
||||
self.assertFalse(result["ok"])
|
||||
|
||||
def test_prd_example_and_whitespace_id_never_submit(self) -> None:
|
||||
cases = [
|
||||
{"ALLOW_LIVE_SUBMIT": "true", "LIVE_CONFIRM": LIVE_CONFIRM_VALUE},
|
||||
{"ALLOW_LIVE_SUBMIT": "1", "LIVE_CONFIRM": "", "CONFIRM_LIVE": "YES"},
|
||||
{"ALLOW_LIVE_SUBMIT": "1", "LIVE_CONFIRM": LIVE_CONFIRM_VALUE, "STRATEGY_ID": " "},
|
||||
]
|
||||
for overrides in cases:
|
||||
with self.subTest(overrides=overrides):
|
||||
client = FakeClient()
|
||||
env = {"STRATEGY_ID": "sid-test", **overrides}
|
||||
with mock.patch.dict(os.environ, env, clear=True):
|
||||
result = attempt_live_submit({}, client=client)
|
||||
self.assertFalse(result["submitted"])
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(client.calls, [])
|
||||
self.assertTrue(result["reason"])
|
||||
|
||||
def test_injected_fake_client(self) -> None:
|
||||
client = FakeClient()
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"STRATEGY_ID": "sid-1",
|
||||
"ALLOW_LIVE_SUBMIT": "1",
|
||||
"LIVE_CONFIRM": LIVE_CONFIRM_VALUE,
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
result = attempt_live_submit({"modelAddress": "x/y"}, client=client)
|
||||
self.assertTrue(result["submitted"])
|
||||
self.assertEqual(len(client.calls), 1)
|
||||
|
||||
def test_default_attempt_rejects(self) -> None:
|
||||
client = FakeClient()
|
||||
with self._clear_gate_env():
|
||||
result = attempt_live_submit({"modelAddress": "x/y"}, client=client)
|
||||
self.assertFalse(result["submitted"])
|
||||
self.assertEqual(len(client.calls), 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
def test_no_client_no_network(self):
|
||||
with patch('urllib.request.urlopen') as network:
|
||||
self.assertFalse(attempt_live_submit({})['submitted'])
|
||||
network.assert_not_called()
|
||||
|
||||
109
tests/test_live_cycle.py
Normal file
109
tests/test_live_cycle.py
Normal file
@@ -0,0 +1,109 @@
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
from src.live_cycle import run_cycle, _submitted
|
||||
|
||||
|
||||
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'}, clear=True)
|
||||
self.env.start()
|
||||
self.addCleanup(self.env.stop)
|
||||
_submitted.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 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('duplicate', self.logs.getvalue())
|
||||
|
||||
def test_remote_duplicate_and_transport_reservation(self):
|
||||
run_cycle([candidate('remote')], self.client, lambda: (1, 'ok', {'org/remote'}))
|
||||
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):
|
||||
run_cycle([candidate('a')], self.client)
|
||||
self.client.submit.assert_not_called()
|
||||
self.assertIn('missing/invalid_token', 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('gpu_unmapped', self.logs.getvalue())
|
||||
self.assertIn('placeholder', self.logs.getvalue())
|
||||
55
tests/test_live_queue.py
Normal file
55
tests/test_live_queue.py
Normal file
@@ -0,0 +1,55 @@
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
from src.live_queue import read_queue
|
||||
|
||||
|
||||
class TestLiveQueue(unittest.TestCase):
|
||||
def response(self, body):
|
||||
r = MagicMock()
|
||||
r.__enter__.return_value = r
|
||||
r.status = 200
|
||||
r.read.return_value = json.dumps(body).encode()
|
||||
return r
|
||||
|
||||
def test_pagination_and_status(self):
|
||||
pages = [self.response({'code': 0, 'data': {'total': 3, 'records': [
|
||||
{'status': 'waiting', 'modelAddress': 'org/a'}, {'status': 'success', 'verifyResult': -1}]}}),
|
||||
self.response({'code': 0, 'data': {'total': 3, 'records': [{'status': '等待中'}]}})]
|
||||
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True), patch('urllib.request.urlopen', side_effect=pages) as call:
|
||||
count, why, models = read_queue()
|
||||
self.assertEqual(count, 2)
|
||||
self.assertIn('org/a', models)
|
||||
self.assertIn('current=2', call.call_args.args[0].full_url)
|
||||
|
||||
def test_fail_closed(self):
|
||||
for body in ({'code': 40100}, {'code': 40101}, {'code': 0, 'data': {'total': 1, 'records': []}}, []):
|
||||
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True), patch('urllib.request.urlopen', return_value=self.response(body)):
|
||||
self.assertIsNone(read_queue()[0])
|
||||
with patch.dict(os.environ, {}, clear=True), patch('urllib.request.urlopen') as call:
|
||||
self.assertIsNone(read_queue()[0])
|
||||
call.assert_not_called()
|
||||
|
||||
def test_page_limit_and_transport(self):
|
||||
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True):
|
||||
with patch('urllib.request.urlopen', return_value=self.response({'code': 0, 'data': {'total': 21, 'records': [{'status': 'waiting'}]}})) as call:
|
||||
self.assertIsNone(read_queue()[0])
|
||||
self.assertEqual(call.call_count, 20)
|
||||
with patch('urllib.request.urlopen', side_effect=OSError('synthetic-secret')):
|
||||
self.assertNotIn('synthetic-secret', str(read_queue()))
|
||||
|
||||
def test_invalid_json_http_and_auth_header(self):
|
||||
with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}, clear=True):
|
||||
r = self.response({})
|
||||
r.read.return_value = b'invalid synthetic-secret'
|
||||
with patch('urllib.request.urlopen', return_value=r):
|
||||
self.assertEqual(read_queue()[1], 'malformed_json')
|
||||
r.status = 503
|
||||
with patch('urllib.request.urlopen', return_value=r):
|
||||
self.assertIsNone(read_queue()[0])
|
||||
with patch('urllib.request.urlopen', return_value=self.response({'code': 0, 'data': {'total': 0, 'records': []}})) as call:
|
||||
self.assertEqual(read_queue()[0], 0)
|
||||
req = call.call_args.args[0]
|
||||
self.assertEqual(req.get_header('Xc-token'), 'synthetic-secret')
|
||||
self.assertEqual(req.method, 'GET')
|
||||
@@ -20,6 +20,7 @@ class TestSigterm(unittest.TestCase):
|
||||
port = "19090"
|
||||
env = os.environ.copy()
|
||||
env["STRATEGY_ID"] = "sigterm-test"
|
||||
env["ALLOW_LIVE_SUBMIT"] = "0"
|
||||
env["PORT"] = port
|
||||
env["HOST"] = "127.0.0.1"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user