import argparse import datetime import gc import hashlib import importlib.metadata import json import os import pathlib import resource import sys import time import torch from huggingface_hub import HfApi, snapshot_download p = argparse.ArgumentParser() p.add_argument('kind', choices=['d1', 'laya', 'decider']) p.add_argument('--root', required=True) p.add_argument('--laya-model', choices=['english','typed-decisions','multilingual']) a = p.parse_args() root = pathlib.Path(a.root) root.mkdir(parents=True, exist_ok=True) import fcntl gpu_lock = (pathlib.Path(__file__).parent / 'gpu-test.lock').open('w') fcntl.flock(gpu_lock, fcntl.LOCK_EX) torch.set_num_threads(8) torch.manual_seed(17) repo = {'d1': 'LiquidAI/d1-omni-600M', 'laya': 'convaiinnovations/laya', 'decider': 'perplexity-ai/pplx-decider-v1.1-27b'}[a.kind] metadata_file = pathlib.Path(__file__).parent / 'official_models.json' if metadata_file.exists() and repo in json.loads(metadata_file.read_text()): from types import SimpleNamespace info = SimpleNamespace(**json.loads(metadata_file.read_text())[repo]) else: info = HfApi().model_info(repo) out = { 'model': repo, 'revision': info.sha, 'date_utc': datetime.datetime.now(datetime.timezone.utc).isoformat(), 'environment': {name: importlib.metadata.version(name) for name in ['torch', 'transformers', 'huggingface_hub']}, 'gpu': torch.cuda.get_device_name(0), 'gpu_total_bytes': torch.cuda.get_device_properties(0).total_memory, 'cuda': torch.version.cuda, 'cases': [], } path = root / (a.kind + ('-' + a.laya_model if a.laya_model else '') + '.json') def save(): path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding='utf-8') def used(): import subprocess r = subprocess.run(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,noheader,nounits'], capture_output=True, text=True) return r.stdout.strip() def case(label, call, gpu=True, extra=None): row = {'label': label, **(extra or {})} try: result = call() if gpu: torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() times = [] for _ in range(3): start = time.perf_counter() result = call() if gpu: torch.cuda.synchronize() times.append((time.perf_counter() - start) * 1000) row.update(status='passed', warmup_excluded=True, latency_ms=times, result=result) if 'expected_choice' in row: row['classification_correct'] = result['answers']['team']['choice'] == row['expected_choice'] if 'expected_choices' in row: row['classification_correct'] = [r['answers']['team']['choice'] for r in result] == row['expected_choices'] if gpu: row.update(peak_allocated_bytes=torch.cuda.max_memory_allocated(), peak_reserved_bytes=torch.cuda.max_memory_reserved(), nvidia_smi_used_mib=used()) row['process_max_rss_kib'] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss except Exception as error: row.update(status='failed', error=repr(error)) out['cases'].append(row) save() print('CASE_JSON=' + json.dumps(row, ensure_ascii=False), flush=True) return row question = {'team': {'type': 'choice', 'instructions': 'Choose a team for this request.', 'criteria': {'billing': 'Payments, invoices and refunds', 'technical': 'Software and integration problems'}}} states = ['I was charged twice. Please refund the duplicate payment.', 'The application crashes at startup. Please help me fix it.'] def valid(payload): assert isinstance(payload, dict) and isinstance(payload.get('answers'), dict) answer = payload['answers']['team'] assert answer['choice'] in ('billing', 'technical') assert 0 <= answer['confidence'] <= 1 assert payload['usage']['output_tokens'] == 0 return payload try: torch.cuda.reset_peak_memory_stats() if a.kind in ('d1', 'decider'): existing = root / 'checkpoint' if a.kind == 'd1' and (existing / 'model.safetensors').is_file() and (existing / 'model.safetensors').stat().st_size == 2_348_774_500: with (existing / 'model.safetensors').open('rb') as stream: assert hashlib.file_digest(stream, 'sha256').hexdigest() == '0713bb05270c2685ad106522f4092bceeeb3a93cf79b401f399a712296c911e1' checkpoint = existing else: checkpoint = pathlib.Path(snapshot_download(repo, revision=info.sha, local_dir=root / 'checkpoint', allow_patterns=['*.json', '*.py', '*.safetensors', '*.txt', '*.model', '*.jinja', '*.tiktoken'])) files = [] for file in checkpoint.glob('*.safetensors'): with file.open('rb') as stream: digest = hashlib.file_digest(stream, 'sha256').hexdigest() files.append({'file': file.name, 'bytes': file.stat().st_size, 'sha256': digest}) out['files'] = files save() if a.kind == 'd1': from PIL import Image from transformers import AutoModel import numpy as np model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True, dtype=torch.float16).to('cuda').eval() out['load_peak_allocated_bytes'] = torch.cuda.max_memory_allocated() out['load_peak_reserved_bytes'] = torch.cuda.max_memory_reserved() for i, state in enumerate(states): case(f'cuda_fp16_text_{i}', lambda s=state: valid(model.system_one(s, question)), extra={'dtype': 'float16', 'input': state, 'expected_choice': ['billing', 'technical'][i]}) for repeats in [10, 40, 160, 640]: state = (states[0] + ' ') * repeats case(f'cuda_fp16_text_repeats{repeats}', lambda s=state: valid(model.system_one(s, question)), extra={'dtype': 'float16', 'state_repetitions': repeats, 'expected_choice': 'billing'}) for size in [2, 8]: requests = [(states[i % 2], question) for i in range(size)] case(f'cuda_fp16_batch{size}', lambda r=requests: [valid(x) for x in model.system_one_batch(r)], extra={'dtype': 'float16', 'batch_size': size, 'expected_choices': [['billing', 'technical'][i % 2] for i in range(size)]}) image = Image.new('RGB', (512, 512), (220, 30, 30)) iq = {'team': {'type': 'choice', 'instructions': 'Choose the dominant color in the image.', 'criteria': {'billing': 'Red', 'technical': 'Blue'}}} case('cuda_fp16_image_512', lambda: valid(model.system_one(None, iq, images=[image])), extra={'dtype': 'float16', 'fixture': 'synthetic solid red RGB image, 512x512', 'expected_choice': 'billing'}) audio = np.zeros(16000, dtype=np.int16) case('cuda_fp16_audio_1s', lambda: valid(model.system_one('Classify this voice clip.', question, audio=audio)), extra={'dtype': 'float16', 'fixture': 'synthetic silence, 16000 Hz mono, 1 second; API check only'}) try: model.system_one(None, question, images=[image], audio=audio) out['media_exclusivity'] = {'status': 'failed', 'error': 'combined image+audio was accepted'} except ValueError as error: out['media_exclusivity'] = {'status': 'passed', 'error': str(error)} del model gc.collect() torch.cuda.empty_cache() model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True, dtype=torch.float32).to('cpu').eval() for i, state in enumerate(states): case(f'cpu_fp32_text_{i}', lambda s=state: valid(model.system_one(s, question)), gpu=False, extra={'dtype': 'float32', 'input': state, 'cpu_threads': 8, 'expected_choice': ['billing', 'technical'][i]}) elif a.kind == 'laya': from laya import Router out['environment']['laya'] = importlib.metadata.version('laya') local = root / 'checkpoint' models = {name: str(local / sub) for name, sub in [('english', ''), ('typed-decisions', 'typed-decisions'), ('multilingual', 'multilingual')]} if local.exists() else None router = Router(models=models, device='cuda', default='english', preload=False, max_loaded=1, revision=info.sha) if local.exists(): out['files'] = [] for file in local.rglob('*.safetensors'): with file.open('rb') as stream: digest = hashlib.file_digest(stream, 'sha256').hexdigest() expected = next(f for f in json.loads(metadata_file.read_text())[repo]['siblings'] if f['rfilename'] == file.relative_to(local).as_posix()) assert digest == expected['lfs']['sha256'] out['files'].append({'file': file.relative_to(local).as_posix(), 'bytes': file.stat().st_size, 'sha256': digest}) for name, length in [('english', 128), ('english', 512), ('typed-decisions', 1024), ('multilingual', 512)]: if a.laya_model and name != a.laya_model: continue agent = router.load(name) gc.collect() torch.cuda.empty_cache() runtime = {'weight_dtypes': sorted({str(p.dtype) for p in agent.model.parameters()}), 'amp_dtype': str(agent.dtype), 'amp_enabled': agent.amp_enabled, 'actual_device': str(agent.device), 'resident_models': router.loaded} assert str(agent.device).startswith('cuda') state = (states[0] + ' ') * (length // 12) case(f'cuda_{name}_limit{length}', lambda s=state, n=name, l=length: valid(router.predict(s, question, model=n, max_len=l)), extra={'requested_model': name, 'max_len': length, 'input': state, 'expected_choice': 'billing', **runtime}) else: sys.path.insert(0, str(checkpoint / 'source' / 'src')) from autojev.model import DecisionModel, answer model = DecisionModel(checkpoint, device='cuda', max_image_pixels=512 * 512).eval() out['load_peak_allocated_bytes'] = torch.cuda.max_memory_allocated() out['load_peak_reserved_bytes'] = torch.cuda.max_memory_reserved() q = question['team'] def predict(state, images=None): row = {'state': state, 'question': q} if images is not None: row['images'] = images probabilities = model.predict([row])[0] assert len(probabilities) == 2 and all(0 <= x <= 1 for x in probabilities) assert abs(sum(probabilities) - 1) < 1e-4 return {'answer': answer(q, probabilities), 'probabilities': probabilities} for length in [128, 512, 2048]: state = (states[0] + ' ') * (length // 12) case(f'cuda_bf16_text_approx{length}', lambda s=state: predict(s), extra={'dtype': 'bfloat16', 'input': state}) from PIL import Image case('cuda_bf16_image_512', lambda: predict(states[1], [Image.new('RGB', (512, 512), (220, 30, 30))]), extra={'dtype': 'bfloat16', 'fixture': 'synthetic solid red RGB image, 512x512'}) out['status'] = 'completed_with_failures' if any(row['status'] == 'failed' for row in out['cases']) else 'completed' except Exception as error: import traceback out['status'] = 'failed' out['error'] = repr(error) out['traceback'] = traceback.format_exc() print(out['traceback'], flush=True) finally: save() print('RESULT_JSON=' + json.dumps(out, ensure_ascii=False), flush=True)