55 lines
4.3 KiB
Python
55 lines
4.3 KiB
Python
import hashlib
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from capacity import assess,overview,GIB
|
|
from catalog import Catalog
|
|
from profiles import Profiles
|
|
|
|
class ManagementTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp=tempfile.TemporaryDirectory();self.root=Path(self.tmp.name);self.model_id='a'*64
|
|
target=self.root/'models'/self.model_id;target.mkdir(parents=True);(target/'model.gguf').write_bytes(b'GGUFtest')
|
|
self.entry=dict(repo='test/Qwen-Image-2.1',file='test.gguf',kind='image',size=8,revision='b'*40,sha256=hashlib.sha256(b'GGUFtest').hexdigest(),downloaded_at=1)
|
|
(target/'entry.json').write_text(json.dumps(self.entry));self.catalog=Catalog(self.root/'models');self.profiles=Profiles(self.root/'profiles.json',self.catalog)
|
|
def tearDown(self):self.tmp.cleanup()
|
|
def request(self):return dict(id=None,revision=0,name='image-test',kind='image',model_id=self.model_id,parameters=dict(width=1024,height=1024,steps=25,seed=-1,guidance=1))
|
|
def test_profiles_persist_and_update_conflicts(self):
|
|
saved=self.profiles.save(self.request());other=Profiles(self.root/'profiles.json',self.catalog)
|
|
row=other.status()['profiles'][0];self.assertEqual(row['parameters']['width'],1024);self.assertFalse(row['runnable']);self.assertIn('Textencoder',row['blockers'][0])
|
|
change=self.request();change.update(id=saved['id'],revision=saved['revision']);change['parameters']['width']=512;other.save(change)
|
|
with self.assertRaises(ValueError):other.save(change)
|
|
self.assertEqual(other.status()['profiles'][0]['parameters']['width'],512)
|
|
def test_profile_rejects_invalid_model_kind_parameters(self):
|
|
for change in [dict(model_id='../secrets'),dict(kind='chat'),dict(name='bad name')]:
|
|
with self.subTest(change=change),self.assertRaises(ValueError):self.profiles.save(dict(self.request(),**change))
|
|
for key,value in [('width',1000),('steps',True),('guidance',float('nan')),('seed',-2)]:
|
|
req=self.request();req['parameters'][key]=value
|
|
with self.subTest(key=key),self.assertRaises(ValueError):self.profiles.save(req)
|
|
(self.root/'models'/self.model_id/'model.gguf').unlink()
|
|
with self.assertRaises(ValueError):self.profiles.save(self.request())
|
|
def test_clear_history_keeps_models_and_survives_restart(self):
|
|
old=self.catalog.status()['downloads'][0];self.catalog.dismiss(old['id'])
|
|
restart=Catalog(self.root/'models');self.assertEqual(restart.status()['downloads'],[]);self.assertEqual(len(restart.status()['entries']),1)
|
|
self.assertEqual(restart.entry(self.model_id)['size'],8)
|
|
def test_running_download_cannot_be_dismissed_and_recovers(self):
|
|
job=dict(id='active',state='downloading',kind='image',repo='test/model',file='a.gguf',bytes=3,total=8)
|
|
self.catalog.history.append(job);self.catalog._save_history()
|
|
with self.assertRaises(ValueError):self.catalog.dismiss('active')
|
|
restarted=Catalog(self.root/'models');self.assertEqual(restarted.status()['downloads'][0]['state'],'interrupted')
|
|
def test_capacity_does_not_sum_gpu_vram_or_claim_runtime_fit(self):
|
|
hw=dict(gpus=[dict(name='A',total_mib=8192,used_mib=4096),dict(name='B',total_mib=8192,used_mib=0)],ram={})
|
|
result=assess(10*GIB,hw,'model.gguf');self.assertFalse(any(x['weights_fit_total'] for x in result['gpus']));self.assertEqual(result['scope'],'weights_only')
|
|
result=assess(6*GIB,hw,'model.gguf');self.assertTrue(result['weights_fit_now']);self.assertFalse(result['gpus'][0]['weights_fit_now'])
|
|
self.assertIn('Gesamtbedarf offen',result['label']);self.assertEqual(assess(1,hw,'config.json')['scope'],'configuration')
|
|
self.assertIsNone(assess(6*GIB,{'gpus':[]})['weights_fit_now'])
|
|
def test_size_range_uses_weight_files_and_labels_scope(self):
|
|
result=overview([dict(name='config.json',size=1),dict(name='a.gguf',size=4*GIB),dict(name='b.gguf',size=8*GIB)],{'gpus':[]})
|
|
self.assertEqual(result['min_bytes'],4*GIB);self.assertEqual(result['max_bytes'],8*GIB);self.assertIn('nicht Gesamtpaket',result['size_note'])
|
|
|
|
def test_variant_sizes_exclude_companions_and_sum_shards(self):
|
|
files=[dict(name='vae/vae.safetensors',size=1),dict(name='model-lora.safetensors',size=2),dict(name='model-00001-of-00002.gguf',size=3*GIB),dict(name='model-00002-of-00002.gguf',size=2*GIB)]
|
|
result=overview(files,{'gpus':[]});self.assertEqual(result['min_bytes'],5*GIB)
|
|
self.assertEqual(overview(files[:-1],{'gpus':[]})['scope'],'unknown')
|