def main()

in src/straggler_healthcheck/pp_benchmark_analysis_runner.py [0:0]

• 40 lines of code
• 6 McCabe index (conditional complexity)

def main(argv: Sequence[str]):
  if len(argv) > 1:
    raise app.UsageError("Too many command-line arguments.")
  data_dir = _DATA_DIR.value
  output_dir = _OUTPUT_DIR.value
  straggler_threshold_ms = _STRAGGLER_THRESHOLD_MS.value
  interesting_event_offset = _INTERESTING_EVENT_OFFSET.value

  experiments = set(
      pathlib.Path(filename).stem.split("-")[1]
      for filename in pathlib.Path(data_dir).glob("*.textproto")
  )
  if not experiments:
    print("No experiments found in dir: %s" % data_dir)
    return

  for experiment in experiments:
    experiment_dir = pathlib.Path(output_dir, experiment)
    experiment_dir.mkdir(parents=True, exist_ok=True)
    experiment_data = pp_benchmark_analysis.read_experiment_data(
        data_dir, experiment
    )
    heatmap_data = pp_benchmark_analysis.preprocess_experiment_data(
        benchmark_results=experiment_data,
        experiment=experiment,
        straggler_threshold_ms=straggler_threshold_ms,
        interesting_event_offset=interesting_event_offset,
    )
    if heatmap_data:
      output_file = pathlib.Path(experiment_dir, _IMAGE_FILE_NAME.value)
      _ = pp_benchmark_analysis.plot_straggler_heatmap(
          heatmap_data=heatmap_data,
          output_file=str(output_file),
          heatmap_metadata=pp_benchmark_analysis.HeatmapMetadata(
              min_width=_MINIMUM_CHART_WIDTH.value,
              max_width=_MAXIMUM_CHART_WIDTH.value,
              min_height=_MINIMUM_CHART_HEIGHT.value,
              max_height=_MAXIMUM_CHART_HEIGHT.value,
          ),
      )
    else:
      print(f"No stragglers found for experiment {experiment}.")