Coverage for dantinox / plotting.py: 19%

1""" 

2Plot generation for DantinoX benchmark results. 

3 

4Delegates to the four plot modules bundled in ``dantinox/plots/``, producing 

5all 16 figures from a benchmark CSV. The public API is the :class:`Plotter` 

6class; the four script modules can also be called directly for advanced use. 

7 

8Usage:: 

9 

10 from dantinox import Plotter 

11 

12 Plotter("benchmark_results.csv").run() 

13 Plotter("benchmark_results.csv").run(groups=["perf", "3d"]) 

14""" 

15 

16from __future__ import annotations 

17 

18import importlib 

19import logging 

20import os 

21import types 

22 

23from dantinox.exceptions import PlotError 

24 

25log = logging.getLogger(__name__) 

26 

27# group name → (module, list of figure functions) 

28_PLOT_GROUPS: dict[str, tuple[str, list[str]]] = { 

29 "perf": ( 

30 "dantinox.plots.plot_perf", 

31 ["fig1_cache_breakdown", "fig2_seqlen_throughput", 

32 "fig3_flops_vs_cache", "fig4_batch_throughput", "fig5_prefill"], 

33 ), 

34 "insights": ( 

35 "dantinox.plots.plot_insights", 

36 ["fig1_pareto", "fig2_serving", "fig3_mla_dial"], 

37 ), 

38 "3d": ( 

39 "dantinox.plots.plot_3d", 

40 ["fig1_cache_surface", "fig2_quality_cube", 

41 "fig3_efficiency_cube", "fig4_serving_surface"], 

42 ), 

43 "3d_dkv": ( 

44 "dantinox.plots.plot_3d_dkv", 

45 ["fig5_dkv_cache_seqlen", "fig6_kv_decoupling", 

46 "fig7_mla_quality", "fig8_dkv_numblocks"], 

47 ), 

48} 

49 

50ALL_GROUPS: list[str] = list(_PLOT_GROUPS.keys()) 

51 

52 

53def _run_group( 

54 group: str, 

55 in_csv: str, 

56 out_dir: str, 

57 batch_csv: str | None, 

58) -> list[str]: 

59 module_name, fig_fns = _PLOT_GROUPS[group] 

60 try: 

61 mod: types.ModuleType = importlib.import_module(module_name) 

62 except ImportError as exc: 

63 raise PlotError(f"Cannot import {module_name}: {exc}") from exc 

64 

65 # Temporarily patch module-level path constants so the scripts write 

66 # to the caller's out_dir and read from the caller's CSV. 

67 orig_in = getattr(mod, "IN_CSV", None) 

68 orig_out = getattr(mod, "OUT_DIR", None) 

69 orig_batch = getattr(mod, "BATCH_CSV", None) 

70 mod.IN_CSV = in_csv # type: ignore[attr-defined] 

71 mod.OUT_DIR = out_dir # type: ignore[attr-defined] 

72 if hasattr(mod, "BATCH_CSV") and batch_csv: 

73 mod.BATCH_CSV = batch_csv # type: ignore[attr-defined] 

74 

75 saved: list[str] = [] 

76 try: 

77 df = mod.load() 

78 for fn_name in fig_fns: 

79 fn = getattr(mod, fn_name, None) 

80 if fn is None: 

81 log.warning(" %s not found in %s — skipped", fn_name, module_name) 

82 continue 

83 if fn_name == "fig4_batch_throughput": 

84 bdf = mod.load_batch() if hasattr(mod, "load_batch") else None 

85 if bdf is not None and not bdf.empty: 

86 fn(bdf) 

87 else: 

88 getattr(mod, "fig4_missing", lambda: None)() 

89 else: 

90 fn(df) 

91 saved.append(fn_name) 

92 log.debug(" %s — done", fn_name) 

93 finally: 

94 if orig_in is not None: 

95 mod.IN_CSV = orig_in # type: ignore[attr-defined] 

96 if orig_out is not None: 

97 mod.OUT_DIR = orig_out # type: ignore[attr-defined] 

98 if orig_batch is not None: 

99 mod.BATCH_CSV = orig_batch # type: ignore[attr-defined] 

100 

101 return saved 

102 

103 

104class Plotter: 

105 """ 

106 Generates all DantinoX benchmark plots from a CSV produced by 

107 :meth:`~dantinox.BenchmarkRunner.run`. 

108 

109 Runs the four bundled plot modules (``perf``, ``insights``, ``3d``, 

110 ``3d_dkv``) and writes 16 PNG files to *out_dir*. 

111 

112 Parameters 

113 ---------- 

114 in_csv : str 

115 Path to ``benchmark_results.csv``. 

116 out_dir : str 

117 Directory where PNGs are written (created if absent). 

118 batch_csv : str, optional 

119 Path to ``batch_sweep_results.csv`` for the batch-throughput plot. 

120 If omitted, that figure is replaced with a placeholder. 

121 

122 Raises 

123 ------ 

124 PlotError 

125 If the CSV is missing or a group name is invalid. 

126 

127 Examples 

128 -------- 

129 >>> from dantinox import BenchmarkRunner, Plotter 

130 >>> BenchmarkRunner("runs").run(out_csv="benchmark_results.csv") 

131 >>> Plotter("benchmark_results.csv").run() 

132 """ 

133 

134 def __init__( 

135 self, 

136 in_csv: str = "benchmark_results.csv", 

137 out_dir: str = "plots", 

138 *, 

139 batch_csv: str | None = None, 

140 ) -> None: 

141 self.in_csv = in_csv 

142 self.out_dir = out_dir 

143 self.batch_csv = batch_csv 

144 

145 def __repr__(self) -> str: 

146 return f"Plotter(in_csv={self.in_csv!r}, out_dir={self.out_dir!r})" 

147 

148 def run(self, groups: list[str] | None = None) -> dict[str, list[str]]: 

149 """ 

150 Generate plots and save them as PNGs. 

151 

152 Parameters 

153 ---------- 

154 groups : list[str], optional 

155 Subset of ``["perf", "insights", "3d", "3d_dkv"]``. 

156 Generates all four if omitted. 

157 

158 Returns 

159 ------- 

160 dict[str, list[str]] 

161 Mapping of group name → list of figure function names that ran. 

162 

163 Raises 

164 ------ 

165 PlotError 

166 If the benchmark CSV is not found or a group name is invalid. 

167 """ 

168 if not os.path.exists(self.in_csv): 

169 raise PlotError( 

170 f"Benchmark CSV not found: {self.in_csv}\n" 

171 "Run BenchmarkRunner.run(out_csv='benchmark_results.csv') first." 

172 ) 

173 

174 os.makedirs(self.out_dir, exist_ok=True) 

175 selected = list(groups) if groups else ALL_GROUPS 

176 unknown = [g for g in selected if g not in _PLOT_GROUPS] 

177 if unknown: 

178 raise PlotError( 

179 f"Unknown plot group(s): {unknown}. Valid groups: {ALL_GROUPS}" 

180 ) 

181 

182 results: dict[str, list[str]] = {} 

183 for group in selected: 

184 log.info("[%s] generating plots…", group) 

185 try: 

186 done = _run_group(group, self.in_csv, self.out_dir, self.batch_csv) 

187 results[group] = done 

188 log.info(" %d figures written to %s/", len(done), self.out_dir) 

189 except PlotError: 

190 raise 

191 except Exception as exc: 

192 log.error(" group '%s' failed: %s", group, exc) 

193 results[group] = [] 

194 

195 return results