Files
Athena-Deck/test_model_management.py
T

108 lines
8.7 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_chat_medium_and_legacy_parameters(self):
self.entry['kind']='chat';(self.root/'models'/self.model_id/'entry.json').write_text(json.dumps(self.entry))
params=dict(context=160000,slots=2,threads=6,batch=2048,ubatch=256)
req=dict(self.request(),kind='chat',name='qwen-medium',parameters=params)
saved=self.profiles.save(req)
self.assertEqual(saved['parameters']['gpu_devices'],[])
devices=['GPU-8ad38c6c-5a01-9d8e-1dfa-ed662ad78fbe','GPU-4834d9d7-5b61-3004-1fb3-4ae49d482d4b']
params.update(gpu_devices=devices,split_mode='layer',tensor_split=[85,15],temperature=1.0,top_p=.95,top_k=20)
req.update(id=saved['id'],revision=1)
self.profiles.save(req)
restarted=Profiles(self.root/'profiles.json',self.catalog)
self.assertEqual(restarted.status()['profiles'][0]['parameters'],dict(params,mtp=False,mtp_tokens=2,mtp_min_p=.05,gpu_offload='auto',vision_projector=None,vision_device='cpu'))
for change in [dict(gpu_offload='unknown'),dict(mtp='yes'),dict(mtp_tokens=0),dict(mtp_tokens=True),dict(mtp_min_p=1.1),dict(gpu_devices=devices*2),dict(tensor_split=[85]),dict(split_mode='none'),dict(top_p=1.1),dict(temperature=float('nan')),dict(top_k=True),dict(tensor_split=[0,100]),dict(gpu_devices=['CUDA0','CUDA1'])]:
with self.subTest(change=change),self.assertRaises(ValueError):self.profiles.save(dict(req,revision=2,parameters=dict(params,**change)))
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','vision_projector')]:
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'])