| import contextlib |
| import gc |
| import http.client |
| import io |
| import json |
| from pathlib import Path |
| import tempfile |
| import threading |
| import unittest |
| import weakref |
|
|
| import playground |
|
|
|
|
| class PlaygroundTests(unittest.TestCase): |
| def setUp(self): |
| temporary = tempfile.TemporaryDirectory() |
| self.addCleanup(temporary.cleanup) |
| self.root = Path(temporary.name) |
| self.rows = [] |
| for index, name in enumerate(('a', 'b')): |
| checkpoint = self.root / (name + '.pt') |
| checkpoint.write_bytes(('fake-' + name).encode()) |
| self.rows.append({'id': name, 'label': name.upper(), 'description': 'Fixture model', |
| 'checkpoint': str(checkpoint), 'checkpoint_step': index, 'calibrated': bool(index), |
| 'sha256': playground.checksum(checkpoint)}) |
| self.config = self.root / 'models.json' |
| self.config.write_text(json.dumps({'models': self.rows, 'default_model': 'a'})) |
| self.catalog, self.default = playground.load_catalog(self.config) |
| self.request = {'model': 'a', 'text': 'Context', 'question': 'Question?', 'options': [' Yes ', 'No']} |
|
|
| def service(self, blocking=None): |
| events, references, payloads = [], [], [] |
| class Backend: |
| def __init__(self, checkpoint, device, max_tokens): |
| self.id = Path(checkpoint).stem |
| self.calibrated, self.max_tokens = self.id == 'b', max_tokens |
| events.append(('load', self.id, threading.get_ident())) |
| references.append(weakref.ref(self)) |
| if blocking is not None: |
| blocking[0].set() |
| if not blocking[1].wait(5): |
| raise RuntimeError('Test synchronization timed out') |
|
|
| def __call__(self, payload): |
| events.append(('score', self.id, threading.get_ident())) |
| payloads.append(payload) |
| if payload['state'] == 'too long': |
| raise ValueError('Input has 1025 tokens; limit is 1024; no truncation') |
| if payload['state'] == 'private-error': |
| raise RuntimeError('/private/checkpoint/secret.pt prompt text') |
| options = payload['questions']['answer']['criteria'] |
| return {'answers': {'answer': {'type': 'choice', 'choice': next(iter(options)), 'confidence': 0.2, 'probabilities': { |
| option: probability for option, probability in zip(options, (0.75, 0.25))}}}, |
| 'usage': {'input_tokens': 42}} |
|
|
| def __del__(self): |
| events.append(('destroy', self.id, threading.get_ident())) |
|
|
| def cleanup(device): |
| gc.collect() |
| events.append(('cleanup', None, threading.get_ident())) |
| service = playground.ModelService(self.catalog, self.default, device='cpu', backend_factory=Backend, |
| preflight=lambda device: events.append(('preflight', None, threading.get_ident())), cleanup=cleanup) |
| self.addCleanup(service.close) |
| return service, events, references, payloads |
|
|
| def test_catalog_allowlist_frozen_hash_and_public_metadata(self): |
| service, _, _, _ = self.service() |
| public = service.models() |
| self.assertEqual(public['default_model'], 'a') |
| self.assertNotIn('checkpoint', public['models'][0]) |
| self.assertNotIn('sha256', public['models'][0]) |
| with self.assertRaises(playground.RequestError) as error: |
| service.score({**self.request, 'model': '/tmp/arbitrary.pt'}) |
| self.assertEqual(error.exception.code, 'unknown_model') |
| (self.root / 'a.pt').write_bytes(b'changed after startup') |
| with contextlib.redirect_stdout(io.StringIO()), self.assertRaises(playground.RequestError) as error: |
| service.score(self.request) |
| self.assertEqual(error.exception.code, 'model_unavailable') |
| self.assertIsNone(service.backend) |
| self.assertEqual(service.status()['status'], 'error') |
|
|
| def test_validation_preserves_options_and_reports_token_limit(self): |
| service, _, _, payloads = self.service() |
| for changes in ({'options': ['x']}, {'options': ['x', 'x']}, {'options': ['x', ' ']}, |
| {'options': ['x', 1]}, {'text': ''}, {'question': None}, {'checkpoint': '/tmp/other'}): |
| with self.assertRaises(playground.RequestError): |
| service.score({**self.request, **changes}) |
| result = service.score(self.request) |
| self.assertEqual(result['probabilities'], [{'option': ' Yes ', 'probability': 0.75}, |
| {'option': 'No', 'probability': 0.25}]) |
| self.assertEqual(payloads[0]['questions']['answer']['criteria'], {' Yes ': None, 'No': None}) |
| self.assertEqual(result['input_tokens'], 42) |
| self.assertFalse(result['model']['calibrated']) |
| with self.assertRaises(playground.RequestError) as error: |
| service.score({**self.request, 'text': 'too long'}) |
| self.assertEqual((error.exception.status, error.exception.code), (422, 'token_limit')) |
| self.assertIn('1025', error.exception.message) |
| capture = io.StringIO() |
| with contextlib.redirect_stdout(capture), self.assertRaises(playground.RequestError) as error: |
| service.score({**self.request, 'text': 'private-error'}) |
| self.assertNotIn('/private', error.exception.message + capture.getvalue()) |
| self.assertNotIn('prompt text', capture.getvalue()) |
| self.assertIn('RuntimeError', capture.getvalue()) |
|
|
| def test_switch_frees_old_model_and_all_backend_work_uses_one_thread(self): |
| service, events, references, _ = self.service() |
| service.score(self.request) |
| service.score({**self.request, 'model': 'b'}) |
| self.assertIsNone(references[0]()) |
| self.assertIsNotNone(references[1]()) |
| self.assertEqual(len({event[2] for event in events}), 1) |
| self.assertNotEqual(events[0][2], threading.get_ident()) |
| destroy = next(index for index, event in enumerate(events) if event[:2] == ('destroy', 'a')) |
| load = next(index for index, event in enumerate(events) if event[:2] == ('load', 'b')) |
| self.assertLess(destroy, load) |
| self.assertTrue(any(event[0] == 'cleanup' for event in events[destroy + 1:load])) |
| self.assertEqual(service.status(), {'status': 'ready', 'loaded_model_id': 'b'}) |
|
|
| def test_concurrent_score_is_busy_and_status_remains_responsive(self): |
| entered, release = threading.Event(), threading.Event() |
| service, _, _, _ = self.service(blocking=(entered, release)) |
| results = [] |
| worker = threading.Thread(target=lambda: results.append(service.score(self.request))) |
| worker.start() |
| try: |
| self.assertTrue(entered.wait(2)) |
| self.assertEqual(service.status()['status'], 'loading') |
| with self.assertRaises(playground.RequestError) as error: |
| service.score(self.request) |
| self.assertEqual((error.exception.status, error.exception.code), (503, 'busy')) |
| finally: |
| release.set() |
| worker.join(5) |
| self.assertFalse(worker.is_alive()) |
| self.assertEqual(len(results), 1) |
|
|
| def test_calibration_metadata_must_match_actual_backend(self): |
| self.catalog['a']['calibrated'] = True |
| service, _, references, _ = self.service() |
| with contextlib.redirect_stdout(io.StringIO()), self.assertRaises(playground.RequestError): |
| service.score(self.request) |
| self.assertIsNone(service.backend) |
| self.assertIsNone(references[0]()) |
|
|
| def test_retained_scoring_error_cannot_hold_old_model_after_switch(self): |
| service, _, references, _ = self.service() |
| with contextlib.redirect_stdout(io.StringIO()), self.assertRaises(playground.RequestError) as caught: |
| service.score({**self.request, 'text': 'private-error'}) |
| retained_error = caught.exception |
| service.score({**self.request, 'model': 'b'}) |
| self.assertIsNotNone(retained_error) |
| self.assertIsNone(references[0]()) |
|
|
| def test_http_assets_origin_limits_and_json_errors(self): |
| service, _, _, _ = self.service() |
| for name in ('playground.html', 'playground.css', 'playground.js'): |
| (self.root / name).write_text('fixture ' + name) |
| server = playground.create_server(service, port=0, web_root=self.root) |
| worker = threading.Thread(target=server.serve_forever, daemon=True) |
| worker.start() |
| self.addCleanup(worker.join, 2) |
| self.addCleanup(server.server_close) |
| self.addCleanup(server.shutdown) |
| def request(method, path, body=None, headers=None): |
| connection = http.client.HTTPConnection('127.0.0.1', server.server_port, timeout=3) |
| connection.request(method, path, body=body, headers=headers or {}) |
| response = connection.getresponse() |
| content = response.read() |
| result = response.status, dict(response.getheaders()), content |
| connection.close() |
| return result |
| status, headers, content = request('GET', '/') |
| self.assertEqual(status, 200) |
| self.assertEqual(headers['X-Content-Type-Options'], 'nosniff') |
| self.assertIn("frame-ancestors 'none'", headers['Content-Security-Policy']) |
| self.assertIn(b'playground.html', content) |
| self.assertEqual(request('GET', '/../../models.json')[0], 404) |
| self.assertEqual(request('GET', '/', headers={'Host': 'attacker.example'})[0], 403) |
| payload = json.dumps(self.request) |
| matching = {'Host': 'localhost:9999', 'Origin': 'http://localhost:9999', 'Content-Type': 'application/json'} |
| self.assertEqual(request('POST', '/api/score', payload, matching)[0], 200) |
| self.assertEqual(request('POST', '/api/score', payload, {**matching, 'Origin': 'http://evil.example'})[0], 403) |
| self.assertEqual(request('POST', '/api/score', payload, {'Content-Type': 'text/plain'})[0], 415) |
| self.assertEqual(request('POST', '/api/score', '{', {'Content-Type': 'application/json'})[0], 400) |
| self.assertEqual(request('POST', '/api/score', b'', {'Content-Type': 'application/json', |
| 'Content-Length': str(playground.MAX_BODY_BYTES + 1)})[0], 413) |
| status, headers, content = request('PUT', '/api/score') |
| self.assertEqual(status, 501) |
| self.assertIn('application/json', headers['Content-Type']) |
| self.assertIn('error', json.loads(content)) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
