Coverage for tests/simulator/test_simulator_types.py: 100%

1from __future__ import annotations 

2 

3import sys 

4import os 

5 

6# Add current path 

7sys.path.append(os.getcwd()) 

8 

9from tests.test_utils import temp_sys_path 

10 

11with temp_sys_path("simulator", "streamwise"): 

12 from sim_types import Model 

13 from sim_types import GPUType 

14 

15 from sim_types_json import models_to_json 

16 from sim_types_json import workflow_to_json 

17 from sim_types_json import policy_to_json 

18 from sim_types_json import model_list_to_json 

19 

20 from models import GemmaModelAllocation 

21 from models import FluxModelAllocation 

22 

23 from model_provisioner.policies import STREAMWISE_POLICY 

24 

25 from workflows import PODCAST_WORKFLOW 

26 

27 

28def test_serialize_models() -> None: 

29 models = { 

30 GPUType.A100: { 

31 Model.GEMMA: [GemmaModelAllocation( 

32 gpu_type=GPUType.A100, 

33 devices=1, replicas=1)] 

34 }, 

35 GPUType.H200: { 

36 Model.FLUX: [FluxModelAllocation( 

37 gpu_type=GPUType.H200, 

38 devices=2, replicas=1)] 

39 }, 

40 } 

41 

42 models_json = models_to_json(models) 

43 

44 assert models_json == ( 

45 "{" 

46 "'A100': {'gemma': {'devices': 1, 'replicas': 1}}," 

47 "'H200': {'flux': {'devices': 2, 'replicas': 1}}" 

48 "}" 

49 ) 

50 

51 

52def test_serialize_workflow() -> None: 

53 workflow = PODCAST_WORKFLOW 

54 workflow_json = workflow_to_json(workflow) 

55 assert workflow_json == ( 

56 '{' 

57 '"total_video_seconds": 600, ' 

58 '"total_scenes": 43, ' 

59 '"total_frames": {"hf": 18000, "ft": 13800}, ' 

60 '"total_subscenes": 171, ' 

61 '"per_subscene_frames": {"hf": 106, "ft": 81}, ' 

62 '"num_steps": {"flux": 25, "hf": 10, "ft": 10}, ' 

63 '"hf_frames": [36, 72, 108, 144, 324], ' 

64 '"ft_frames": [9, 21, 41, 61, 77], ' 

65 '"frames_per_step_idx": 4, ' 

66 '"target_resolution": "high", ' 

67 '"total_input_tokens": 20480, ' 

68 '"model_work": {' 

69 '"gemma": 1, ' 

70 '"flux": 1, ' 

71 '"hf": 171, ' 

72 '"hf_vae": 18000, ' 

73 '"ft": 171, ' 

74 '"ft_vae": 13800, ' 

75 '"upscaler": 13800, ' 

76 '"others": 1}' 

77 '}' 

78 ) 

79 

80 

81def test_serialize_policy() -> None: 

82 policy = STREAMWISE_POLICY 

83 policy_json = policy_to_json(policy) 

84 assert policy_json == ( 

85 '{' 

86 '"name": "streamwise", ' 

87 '"objective": "Objective.TTFF_COST", ' 

88 '"disaggregation": {"hf": true, "ft": false}, ' 

89 '"use_upscaler": true, ' 

90 '"hardware": ["A100", "H100", "H200", "GB200"]' 

91 '}' 

92 ) 

93 

94 

95def test_serialize_model_list() -> None: 

96 models = [Model.GEMMA, Model.FLUX] 

97 models_json = model_list_to_json(models) 

98 assert models_json == '["gemma", "flux"]'