add minimal model compatibility preflight agent
This commit is contained in:
48
test_main.py
Normal file
48
test_main.py
Normal file
@@ -0,0 +1,48 @@
|
||||
import unittest
|
||||
|
||||
from main import ValidationError, analyze
|
||||
|
||||
|
||||
class AnalyzeTests(unittest.TestCase):
|
||||
def test_feasible_preflight(self):
|
||||
result = analyze(
|
||||
{
|
||||
"model_name": "example/7b-model",
|
||||
"architecture": "transformer",
|
||||
"parameters_b": 7,
|
||||
"dtype": "bf16",
|
||||
"hardware": "国产算力卡",
|
||||
"device_memory_gb": 24,
|
||||
"cards": 1,
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["verdict"], "preflight_feasible")
|
||||
self.assertGreater(result["estimates"]["headroom_gb"], 0)
|
||||
|
||||
def test_insufficient_memory(self):
|
||||
result = analyze(
|
||||
{
|
||||
"model_name": "example/32b-model",
|
||||
"architecture": "transformer",
|
||||
"parameters_b": 32,
|
||||
"dtype": "bf16",
|
||||
"hardware": "国产算力卡",
|
||||
"device_memory_gb": 24,
|
||||
"cards": 2,
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["verdict"], "insufficient_memory")
|
||||
|
||||
def test_missing_information_is_explicit(self):
|
||||
result = analyze({"model_name": "example/unknown-model"})
|
||||
self.assertEqual(result["verdict"], "information_required")
|
||||
self.assertIn("parameters_b", result["missing_fields"])
|
||||
self.assertIn("device_memory_gb", result["missing_fields"])
|
||||
|
||||
def test_invalid_card_count_is_rejected(self):
|
||||
with self.assertRaises(ValidationError):
|
||||
analyze({"cards": 0})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user