Coverage for simulator/data_loading.py: 99%

1""" 

2Module for loading latency and power consumption data from CSV files. 

3""" 

4 

5import pandas as pd 

6 

7from pathlib import Path 

8 

9from sim_types import LatencyData 

10from sim_types import PowerData 

11from sim_types import GPUType 

12from sim_types import LatencyGPUTypeData 

13from sim_types import PowerGPUTypeData 

14from sim_types import QualityLevel 

15 

16from constants import NUM_PIXELS_ORIGINAL_UPSCALER 

17from constants import NUM_PIXELS_ORIGINAL_FT 

18from constants import NUM_PIXELS_ORIGINAL_HF 

19from constants import NUM_PIXELS_ORIGINAL_FLUX 

20from constants import NUM_PIXELS_LOW_FT 

21from constants import NUM_PIXELS_LOW_HF 

22from constants import NUM_PIXELS_LOW_FLUX 

23from constants import NUM_PIXELS_LOW_UPSCALER 

24from constants import NUM_PIXELS_MEDIUM_FT 

25from constants import NUM_PIXELS_MEDIUM_HF 

26from constants import NUM_PIXELS_MEDIUM_UPSCALER 

27from constants import NUM_PIXELS_MEDIUM_FLUX 

28from constants import POWER_GPU_IDLE 

29from constants import POWER_GPU_TDP 

30 

31_DEFAULT_DATA_DIR = Path("data") 

32 

33 

34def load_latency_data( 

35 data_dir: str | Path = _DEFAULT_DATA_DIR, 

36) -> LatencyData: 

37 """ 

38 Load latency and throughput mapping data from CSV files. 

39 

40 Args: 

41 data_dir: The directory where the CSV files are stored. 

42 Returns: 

43 LatencyData: An object containing all loaded latency data. 

44 """ 

45 data_path = Path(data_dir) 

46 

47 data = LatencyData(gpus={}) 

48 for gpu_type in GPUType: 

49 data.gpus[gpu_type] = LatencyGPUTypeData(gpu_type=gpu_type) 

50 

51 # Flux time -> per image generation 

52 csv_flux_path = data_path / f"latency_flux_mapping_{gpu_type.value.lower()}.csv" 

53 df_flux = pd.read_csv(csv_flux_path, comment='#') 

54 data[gpu_type].flux = dict(zip( 

55 df_flux["world_size"], 

56 df_flux["avg_steps_time"])) 

57 

58 # Hunyuan Framepack per step time -> [36, 72, 108, 144, 324] frames generation 

59 csv_hf_path = data_path / f"latency_hf_mapping_{gpu_type.value.lower()}.csv" 

60 df_hf = pd.read_csv(csv_hf_path, comment='#') 

61 data[gpu_type].hf = dict(zip( 

62 df_hf["world_size"], 

63 df_hf["avg_steps_time"])) 

64 

65 # Hunyuan Framepack VAE time -> per inference iteration 

66 # Derived: steps * avg_step_time * vae_pct(vae_time / total_time) 

67 data[gpu_type].hf_vae = dict(zip( 

68 df_hf["world_size"], 

69 df_hf["vae_time"])) 

70 

71 # Fantasy Talking per step time -> [9, 21, 41, 61, 77] frames generation 

72 csv_ft_path = data_path / f"latency_ft_mapping_{gpu_type.value.lower()}.csv" 

73 df_ft = pd.read_csv(csv_ft_path, comment='#') 

74 data[gpu_type].ft = dict(zip( 

75 df_ft["world_size"], 

76 df_ft["avg_steps_time"])) 

77 

78 # Fantasy Talking VAE time -> per inference iteration 

79 # Derived: steps * avg_step_time * vae_pct(vae_time / total_time) 

80 data[gpu_type].ft_vae = dict(zip( 

81 df_ft["world_size"], 

82 df_ft["vae_time"])) 

83 

84 # Upscaler time -> per image frame 

85 csv_upscaler_path = data_path / f"latency_upscaler_{gpu_type.value.lower()}.csv" 

86 df_upscaler = pd.read_csv(csv_upscaler_path, comment='#') 

87 data[gpu_type].upscaler = dict(zip( 

88 df_upscaler['world_size'], 

89 df_upscaler['avg_steps_time'])) 

90 

91 # Gemma time -> first scene and per scene 

92 csv_gemma_path = data_path / f"latency_gemma_{gpu_type.value.lower()}.csv" 

93 df_gemma = pd.read_csv(csv_gemma_path, comment='#') 

94 data[gpu_type].gemma_first_scene = dict(zip( 

95 df_gemma['tp'], 

96 df_gemma['first_scene_time'])) 

97 data[gpu_type].gemma_per_scene = dict(zip( 

98 df_gemma['tp'], 

99 df_gemma['per_scene_time'])) 

100 

101 # Others time -> kokoro and other overheads -> time per scene 

102 csv_others_path = data_path / f"latency_others_{gpu_type.value.lower()}.csv" 

103 df_others = pd.read_csv(csv_others_path, comment='#') 

104 data[gpu_type].others = dict(zip( 

105 df_others['world_size'], 

106 df_others['time'])) 

107 

108 return data 

109 

110 

111def load_power_data( 

112 data_dir: str | Path = _DEFAULT_DATA_DIR 

113) -> PowerData: 

114 """ 

115 Load power consumption data from CSV files. 

116 

117 Args: 

118 data_dir: The directory where the CSV files are stored. 

119 Returns: 

120 PowerData: An object containing all loaded power consumption data. 

121 """ 

122 data_path = Path(data_dir) 

123 

124 data = PowerData(gpus={}) 

125 for gpu_type in GPUType: 

126 data.gpus[gpu_type] = PowerGPUTypeData(gpu_type=gpu_type) 

127 

128 # Flux power profile 

129 power_flux_file_name = data_path / f'power_flux_mapping_{gpu_type.value.lower()}.csv' 

130 power_flux_df = pd.read_csv(power_flux_file_name, comment='#') 

131 data[gpu_type].flux = dict(zip( 

132 power_flux_df['world_size'], 

133 power_flux_df['power_watts'])) 

134 

135 # Hunyuan Framepack 640x400 power profile 

136 power_hf_file_name = data_path / f'power_hf_mapping_{gpu_type.value.lower()}.csv' 

137 power_hf_df = pd.read_csv(power_hf_file_name, comment='#') 

138 data[gpu_type].hf = dict(zip( 

139 power_hf_df['world_size'], 

140 power_hf_df['power_watts'])) 

141 

142 # Hunyuan Framepack 1280x800 power profile 

143 power_hf_file_name_high = data_path / f'power_hf_mapping_{gpu_type.value.lower()}_high.csv' 

144 power_hf_high_df = pd.read_csv(power_hf_file_name_high, comment='#') 

145 data[gpu_type].hf_high = dict(zip( 

146 power_hf_high_df['world_size'], 

147 power_hf_high_df['power_watts'])) 

148 

149 # Hunyuan Framepack VAE power profile 

150 power_hf_vae_file_name = data_path / f'power_hf_vae_{gpu_type.value.lower()}.csv' 

151 power_hf_vae_df = pd.read_csv(power_hf_vae_file_name, comment='#') 

152 data[gpu_type].hf_vae = dict(zip( 

153 power_hf_vae_df['world_size'], 

154 power_hf_vae_df['power_watts'])) 

155 

156 # Hunyuan Framepack VAE high power profile 

157 power_hf_vae_high_file_name = data_path / f'power_hf_vae_{gpu_type.value.lower()}_high.csv' 

158 power_hf_vae_high_df = pd.read_csv(power_hf_vae_high_file_name, comment='#') 

159 data[gpu_type].hf_vae_high = dict(zip( 

160 power_hf_vae_high_df['world_size'], 

161 power_hf_vae_high_df['power_watts'])) 

162 

163 # Fantasy Talking 640x400 power profile 

164 power_ft_file_name = data_path / f'power_ft_mapping_{gpu_type.value.lower()}.csv' 

165 power_ft_df = pd.read_csv(power_ft_file_name, comment='#') 

166 data[gpu_type].ft = dict(zip( 

167 power_ft_df['world_size'], 

168 power_ft_df['power_watts'])) 

169 

170 # Fantasy Talking 1280x800 power profile 

171 power_ft_high_file_name = data_path / f'power_ft_mapping_{gpu_type.value.lower()}_high.csv' 

172 power_ft_high_df = pd.read_csv(power_ft_high_file_name, comment='#') 

173 data[gpu_type].ft_high = dict(zip( 

174 power_ft_high_df['world_size'], 

175 power_ft_high_df['power_watts'])) 

176 

177 # Fantasy Talking VAE mapping 

178 power_ft_vae_file_name = data_path / f'power_ft_vae_mapping_{gpu_type.value.lower()}.csv' 

179 power_ft_vae_df = pd.read_csv(power_ft_vae_file_name, comment='#') 

180 data[gpu_type].ft_vae = dict(zip( 

181 power_ft_vae_df['world_size'], 

182 power_ft_vae_df['power_watts'])) 

183 

184 # Fantasy Talking VAE high mapping 

185 power_ft_vae_high_file_name = data_path / f'power_ft_vae_mapping_{gpu_type.value.lower()}_high.csv' 

186 power_ft_vae_high_df = pd.read_csv(power_ft_vae_high_file_name, comment='#') 

187 data[gpu_type].ft_vae_high = dict(zip( 

188 power_ft_vae_high_df['world_size'], 

189 power_ft_vae_high_df['power_watts'])) 

190 

191 # Upscaler power profile 

192 power_upscaler_file_name = data_path / f'power_upscaler_{gpu_type.value.lower()}.csv' 

193 power_upscaler_df = pd.read_csv(power_upscaler_file_name, comment='#') 

194 data[gpu_type].upscaler = dict(zip( 

195 power_upscaler_df['world_size'], 

196 power_upscaler_df['power_watts'])) 

197 

198 # Gemma power profile 

199 power_gemma_first_scene_file_name = data_path / f'power_gemma_first_scene_{gpu_type.value.lower()}.csv' 

200 power_gemma_per_scene_file_name = data_path / f'power_gemma_per_scene_{gpu_type.value.lower()}.csv' 

201 power_gemma_first_scene_df = pd.read_csv(power_gemma_first_scene_file_name, comment='#') 

202 power_gemma_per_scene_df = pd.read_csv(power_gemma_per_scene_file_name, comment='#') 

203 data[gpu_type].gemma_first_scene = dict(zip( 

204 power_gemma_first_scene_df['world_size'], 

205 power_gemma_first_scene_df['power_watts'] 

206 )) 

207 data[gpu_type].gemma_per_scene = dict(zip( 

208 power_gemma_per_scene_df['world_size'], 

209 power_gemma_per_scene_df['power_watts'] 

210 )) 

211 

212 # Idle and TDP power profiles 

213 for gpu_type in GPUType: 

214 data[gpu_type].idle = POWER_GPU_IDLE[gpu_type] 

215 data[gpu_type].tdp = POWER_GPU_TDP[gpu_type] 

216 

217 return data 

218 

219 

220def load_adaptive_quality_data( 

221 data_dir: str | Path, 

222 level: QualityLevel, 

223) -> LatencyData: 

224 """Load latency data for adaptive quality.""" 

225 assert isinstance(level, QualityLevel) 

226 

227 latency_data = load_latency_data(data_dir=data_dir) 

228 

229 if level == QualityLevel.ORIGINAL or level == QualityLevel.HIGH: 

230 return latency_data 

231 

232 if level == QualityLevel.MEDIUM: 

233 ratio_flux = NUM_PIXELS_MEDIUM_FLUX / NUM_PIXELS_ORIGINAL_FLUX 

234 ratio_hf = NUM_PIXELS_MEDIUM_HF / NUM_PIXELS_ORIGINAL_HF 

235 ratio_hf_vae = NUM_PIXELS_MEDIUM_HF / NUM_PIXELS_ORIGINAL_HF 

236 ratio_ft = NUM_PIXELS_MEDIUM_FT / NUM_PIXELS_ORIGINAL_FT 

237 ratio_ft_vae = NUM_PIXELS_MEDIUM_FT / NUM_PIXELS_ORIGINAL_FT 

238 ratio_upscaler = NUM_PIXELS_MEDIUM_UPSCALER / NUM_PIXELS_ORIGINAL_UPSCALER 

239 for gpu_type in GPUType: 

240 latency_data[gpu_type].flux = { 

241 k: v * ratio_flux 

242 for k, v in latency_data[gpu_type].flux.items() 

243 } 

244 latency_data[gpu_type].hf = { 

245 k: v * ratio_hf 

246 for k, v in latency_data[gpu_type].hf.items() 

247 } 

248 latency_data[gpu_type].hf_vae = { 

249 k: v * ratio_hf_vae 

250 for k, v in latency_data[gpu_type].hf_vae.items() 

251 } 

252 latency_data[gpu_type].ft = { 

253 k: v * ratio_ft 

254 for k, v in latency_data[gpu_type].ft.items() 

255 } 

256 latency_data[gpu_type].ft_vae = { 

257 k: v * ratio_ft_vae 

258 for k, v in latency_data[gpu_type].ft_vae.items() 

259 } 

260 latency_data[gpu_type].upscaler = { 

261 k: v * ratio_upscaler 

262 for k, v in latency_data[gpu_type].upscaler.items() 

263 } 

264 return latency_data 

265 

266 if level == QualityLevel.LOW: 

267 ratio_flux = NUM_PIXELS_LOW_FLUX / NUM_PIXELS_ORIGINAL_FLUX 

268 ratio_hf = NUM_PIXELS_LOW_HF / NUM_PIXELS_ORIGINAL_HF 

269 ratio_hf_vae = NUM_PIXELS_LOW_HF / NUM_PIXELS_ORIGINAL_HF 

270 ratio_ft = NUM_PIXELS_LOW_FT / NUM_PIXELS_ORIGINAL_FT 

271 ratio_ft_vae = NUM_PIXELS_LOW_FT / NUM_PIXELS_ORIGINAL_FT 

272 ratio_upscaler = NUM_PIXELS_LOW_UPSCALER / NUM_PIXELS_ORIGINAL_UPSCALER 

273 for gpu_type in GPUType: 

274 latency_data[gpu_type].flux = { 

275 k: v * ratio_flux 

276 for k, v in latency_data[gpu_type].flux.items() 

277 } 

278 latency_data[gpu_type].hf = { 

279 k: v * ratio_hf 

280 for k, v in latency_data[gpu_type].hf.items() 

281 } 

282 latency_data[gpu_type].hf_vae = { 

283 k: v * ratio_hf_vae 

284 for k, v in latency_data[gpu_type].hf_vae.items() 

285 } 

286 latency_data[gpu_type].ft = { 

287 k: v * ratio_ft 

288 for k, v in latency_data[gpu_type].ft.items() 

289 } 

290 latency_data[gpu_type].ft_vae = { 

291 k: v * ratio_ft_vae 

292 for k, v in latency_data[gpu_type].ft_vae.items() 

293 } 

294 latency_data[gpu_type].upscaler = { 

295 k: v * ratio_upscaler 

296 for k, v in latency_data[gpu_type].upscaler.items() 

297 } 

298 return latency_data 

299 

300 return latency_data