93 lines
7.3 KiB
Python
93 lines
7.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_delete_revision_and_model_preservation(self):
|
|
row=self.profiles.save(self.request())
|
|
with self.assertRaises(ValueError):self.profiles.delete(dict(id=row['id'],revision=0))
|
|
self.assertEqual(len(self.profiles.rows),1)
|
|
self.profiles.delete(dict(id=row['id'],revision=row['revision']))
|
|
self.assertEqual(Profiles(self.root/'profiles.json',self.catalog).rows,[])
|
|
self.assertTrue((self.root/'models'/self.model_id/'model.gguf').exists())
|
|
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')
|
|
def test_component_recipe_assignment_and_preservation(self):
|
|
from profiles import QWEN_REPO,QWEN_COMPONENTS
|
|
from unittest.mock import patch
|
|
main=self.root/'models'/self.model_id/'entry.json';entry=json.loads(main.read_text());entry['repo']=QWEN_REPO;main.write_text(json.dumps(entry))
|
|
ids={};files=[]
|
|
for n,(role,recipe) in enumerate(QWEN_COMPONENTS.items(),1):
|
|
ident=str(n)*64;ids[role]=ident;target=self.root/'models'/ident;target.mkdir();(target/'model.safetensors').write_bytes(b'component')
|
|
f=recipe['files'][0];(target/'entry.json').write_text(json.dumps(dict(repo=QWEN_REPO,file=f,kind='image',size=9,revision='c'*40)))
|
|
files.append(dict(name=f,size=9,sha256=None))
|
|
saved=self.profiles.save(self.request())
|
|
with patch.object(self.catalog,'files',return_value=dict(files=files,revision='c'*40,gated=False,url='https://huggingface.co/'+QWEN_REPO,license='test')):
|
|
result=self.profiles.components(self.model_id)
|
|
self.assertTrue(result['supported']);self.assertEqual(len(result['requirements'][0]['available']),1)
|
|
assigned=self.profiles.assign(dict(id=saved['id'],revision=1,components=ids));self.assertEqual(assigned['revision'],2)
|
|
restarted=Profiles(self.root/'profiles.json',self.catalog);row=restarted.status()['profiles'][0]
|
|
self.assertEqual(row['components'],ids);self.assertEqual(len(row['blockers']),1);self.assertFalse(row['runnable'])
|
|
request=self.request();request.update(id=saved['id'],revision=2);self.assertEqual(restarted.save(request)['components'],ids)
|
|
with self.assertRaises(ValueError):restarted.assign(dict(id=saved['id'],revision=3,components={'vae':ids['text_encoder']}))
|
|
with self.assertRaises(ValueError):restarted.assign(dict(id=saved['id'],revision=2,components=ids))
|
|
request=self.request();request['model_id']=ids['vae']
|
|
with self.assertRaises(ValueError):restarted.save(request)
|
|
def test_auxiliary_files_cannot_be_profile_models(self):
|
|
from catalog import file_role
|
|
for name,role in [('text_encoders/qwen3vl_8b_int8_convrot.safetensors','text_encoder'),('vae/qwen_image_2.1_vae_bf16.safetensors','vae'),('qwen_image_vae.safetensors','vae'),('mmproj-model.gguf','auxiliary')]:
|
|
with self.subTest(name=name):
|
|
self.assertEqual(file_role(name)['role'],role)
|
|
path=self.root/'models'/self.model_id/'entry.json';data=json.loads(path.read_text());data['file']=name;path.write_text(json.dumps(data))
|
|
suffix=Path(name).suffix;(path.parent/('model'+suffix)).write_bytes(b'GGUFtest')
|
|
self.assertFalse(self.catalog.status()['entries'][0]['profile_eligible'])
|
|
with self.assertRaises(ValueError):self.profiles.save(self.request())
|
|
self.assertTrue(file_role('qwen-image-2.1-UC-Q8_0.gguf')['profile_eligible'])
|