Add bounded context ranges and complete GPU split coverage to auto test

test_auto_test.py

@@ -1,8 +1,8 @@
-import tempfile,unittest,json
+import tempfile,unittest,json,time
 from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import Mock,patch
-from auto_test import AutoTests,model_layers,fine_splits
+from auto_test import AutoTests,model_layers,fine_splits,MAX_CANDIDATES,GPU_SPLITS
 from inference import Scheduler,InferenceError
 
 class AutoTestsTests(unittest.TestCase):
@@ -12,19 +12,22 @@ class AutoTestsTests(unittest.TestCase):
   self.gpus=[dict(name='RTX 5080',uuid='first'),dict(name='RTX 3060',uuid='second')]
  def tearDown(self):self.tmp.cleanup()
  def start(self):
-  with patch('auto_test.probe',return_value=self.gpus):self.auto.start(dict(model_id='m',max_context=4096,mtp=False))
+  with patch('auto_test.probe',return_value=self.gpus):self.auto.start(dict(model_id='m',start_context=4096,max_context=4096,mtp=False))
   self.auto.thread.join(3);self.assertFalse(self.auto.thread.is_alive())
  def test_primary_first_results_persist_and_save_explicitly(self):
   self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':40})
-  self.start();job=self.auto.status()['job'];self.assertEqual(job['state'],'complete');self.assertEqual(len(job['results']),6)
-  self.assertTrue(all(r['parameters']['gpu_devices']==['first'] for r in job['results']));self.profiles.save.assert_not_called();self.assertEqual(self.scheduler.active,0)
+  self.start();job=self.auto.status()['job'];self.assertEqual(job['state'],'complete');self.assertEqual(len(job['results']),3*(1+len(GPU_SPLITS)))
+  self.assertEqual(job['completed_contexts'],[4096]);self.assertEqual(job['results'][0]['parameters']['gpu_devices'],['first']);self.profiles.save.assert_not_called();self.assertEqual(self.scheduler.active,0)
+  tested_splits={tuple(r['parameters']['tensor_split']) for r in job['results']}
+  self.assertIn((75,25),tested_splits);self.assertIn((50,50),tested_splits)
   self.auto.save(dict(job_id=job['id'],index=0,name='saved'));self.profiles.save.assert_called_once()
   self.assertEqual(AutoTests(self.auto.path,self.profiles,self.worker,self.scheduler).job['state'],'complete')
  def test_secondary_before_cpu(self):
   def plan(p,cancel):
    if p['parameters']['gpu_offload']=='full':raise InferenceError('does not fit')
   self.worker.plan.side_effect=plan;self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':10});self.start()
-  rows=self.auto.job['results'];first_success=next(i for i,r in enumerate(rows) if r['success']);self.assertEqual(first_success,18);self.assertEqual(rows[first_success]['parameters']['gpu_offload'],'auto')
+  rows=self.auto.job['results'];first_success=next(i for i,r in enumerate(rows) if r['success']);self.assertEqual(first_success,3*(1+len(GPU_SPLITS)));self.assertEqual(rows[first_success]['parameters']['gpu_offload'],'auto')
+  self.assertEqual(rows[0]['failure_stage'],'prediction')
  def test_busy_lease_does_not_stop_another_worker(self):
   with self.scheduler.lease(('other',)):
    self.worker.reset_mock();self.start();self.assertEqual(self.auto.job['state'],'failed');self.worker.stop.assert_not_called()
@@ -32,10 +35,15 @@ class AutoTestsTests(unittest.TestCase):
   self.worker.plan.side_effect=lambda *a,**k:self.auto.cancel.set();self.auto.request=Mock(side_effect=InterruptedError());self.start();self.assertEqual(self.auto.job['state'],'cancelled');self.assertEqual(self.scheduler.active,0);self.profiles.save.assert_not_called()
  def test_validation_and_failed_result_save(self):
   with self.assertRaises(ValueError):self.auto.start(dict(model_id='m',max_context=True,mtp=False))
+  with self.assertRaises(ValueError):self.auto.start(dict(model_id='m',start_context=65536,max_context=8192,mtp=False))
   self.auto.job={'id':'a','model_id':'m','results':[{'success':False}]}
   with self.assertRaises(ValueError):self.auto.save(dict(job_id='a',index=0,name='x'))
  def test_prompt_filled_by_token_count(self):
-  self.auto.request=Mock(side_effect=[{'tokens':list(range(4000))},{'timings':{'prompt_n':3072,'predicted_n':64,'predicted_per_second':12}}]);r=self.auto.benchmark(4096);self.assertEqual(r['prompt_n'],3072);self.assertEqual(len(self.auto.request.call_args.args[1]['prompt']),3072)
+  self.auto.request=Mock(side_effect=[{'tokens':list(range(4000))},{'timings':{'prompt_n':3968,'predicted_n':64,'predicted_per_second':12}}]);r=self.auto.benchmark(4096);self.assertEqual(r['prompt_n'],3968);self.assertEqual(len(self.auto.request.call_args.args[1]['prompt']),3968)
+ def test_budget_exhaustion_is_not_reported_as_complete(self):
+  self.auto.job=dict(id='j',model_id='m',start_context=65536,max_context=65536,mtp=False,attempts=MAX_CANDIDATES,results=[],completed_contexts=[],started_at=time.time(),max_candidates=MAX_CANDIDATES)
+  self.auto.run('first','second')
+  self.assertEqual(self.auto.job['state'],'partial');self.assertEqual(self.auto.job['completed_contexts'],[])
 class FineTests(unittest.TestCase):
  def test_single_layer_steps_between_coarse_candidates(self):
   import math
@@ -55,12 +63,12 @@ class FineTests(unittest.TestCase):
  def test_refinement_integrated_after_first_fit(self):
   fixture=AutoTestsTests();fixture.setUp()
   try:
-   a=fixture.auto;a.job=dict(id='j',model_id='m',max_context=4096,mtp=True,attempts=0,results=[],started_at=0)
+   a=fixture.auto;a.job=dict(id='j',model_id='m',start_context=4096,max_context=4096,mtp=True,attempts=0,results=[],completed_contexts=[],started_at=time.time())
    seen=[]
    def candidate(context,tier,devices,split,ratio,offload,layer_count=None):
     seen.append((context,ratio,layer_count));return [True] if ratio and ratio[0]<=90 else []
    a.test_candidate=candidate;fixture.worker.catalog.root=Path('/unused')
    with patch('auto_test.model_layers',return_value=66):a.run('first','second')
-   self.assertEqual([v[2] for v in seen if v[2] is not None],[61,61])
+   self.assertEqual([v[2] for v in seen if v[2] is not None],[61])
   finally:fixture.tearDown()
 if __name__=='__main__':unittest.main()