Coverage for simulator/sim_types_json.py: 100%

1from __future__ import annotations 

2 

3import json 

4 

5from dataclasses import asdict 

6 

7from sim_types import Model 

8from sim_types import Policy 

9from sim_types import GPUType 

10from sim_types import ModelAllocation 

11from sim_types import WorkflowConfig 

12 

13 

14def models_to_json( 

15 models: dict[GPUType, dict[Model, list[ModelAllocation]]] 

16) -> str: 

17 result = {} 

18 for gpu_type, model_dict in models.items(): 

19 inner_result = {} 

20 for model, allocation_list in model_dict.items(): 

21 for allocation in allocation_list: 

22 alloc_dict = { 

23 'devices': allocation.devices, 

24 'replicas': allocation.replicas, 

25 } 

26 inner_result[model.value] = alloc_dict 

27 result[gpu_type.name] = inner_result 

28 return str(result).replace("}}, '", "}},'") 

29 

30 

31def workflow_to_json(workflow: WorkflowConfig) -> str: 

32 d = asdict(workflow) 

33 # Convert Model enum keys in dict fields to string values 

34 for dict_field in ('total_frames', 'per_subscene_frames', 'num_steps', 'model_work'): 

35 if dict_field in d: 

36 d[dict_field] = { 

37 (k.value if hasattr(k, 'value') else k): v 

38 for k, v in d[dict_field].items() 

39 } 

40 # Convert QualityLevel enum to string value 

41 if 'target_resolution' in d and hasattr(d['target_resolution'], 'value'): 

42 d['target_resolution'] = d['target_resolution'].value 

43 return json.dumps(d) 

44 

45 

46def policy_to_json(policy: Policy) -> str: 

47 result = { 

48 'name': policy.name, 

49 'objective': str(policy.objective), 

50 'disaggregation': {model.value: enabled for model, enabled in policy.disaggregation.items()}, 

51 'use_upscaler': policy.use_upscaler, 

52 'hardware': [gpu.name for gpu in policy.hardware], 

53 } 

54 return json.dumps(result) 

55 

56 

57def model_list_to_json(models: list[Model]) -> str: 

58 return json.dumps(models, default=lambda o: o.value)