From 4ff1cfa6f0e907d41bbe164a20f5e99712390fbf Mon Sep 17 00:00:00 2001 From: Tom White Date: Fri, 16 Feb 2024 12:17:37 +0000 Subject: [PATCH] Change timeline visualization default format to SVG (#384) --- cubed/extensions/timeline.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/cubed/extensions/timeline.py b/cubed/extensions/timeline.py index 59ffd42c..1c0abc89 100644 --- a/cubed/extensions/timeline.py +++ b/cubed/extensions/timeline.py @@ -1,6 +1,7 @@ import os import time from dataclasses import asdict +from typing import Optional import matplotlib.patches as mpatches import numpy as np @@ -15,6 +16,9 @@ class TimelineVisualizationCallback(Callback): + def __init__(self, format: Optional[str] = None) -> None: + self.format = format + def on_compute_start(self, event): self.start_tstamp = time.time() self.stats = [] @@ -25,11 +29,12 @@ def on_task_end(self, event): def on_compute_end(self, event): end_tstamp = time.time() dst = f"history/{event.compute_id}" - create_timeline(self.stats, self.start_tstamp, end_tstamp, dst) + format = self.format + create_timeline(self.stats, self.start_tstamp, end_tstamp, dst, format) # copy of lithops function of the same name, and modified for different field names -def create_timeline(stats, start_tstamp, end_tstamp, dst=None): +def create_timeline(stats, start_tstamp, end_tstamp, dst=None, format=None): stats_df = pd.DataFrame(stats) stats_df = stats_df.sort_values(by=["task_create_tstamp", "name"], ascending=True) @@ -81,13 +86,16 @@ def create_timeline(stats, start_tstamp, end_tstamp, dst=None): ax.grid(False) fig.tight_layout() + if format is None: + format = "svg" + if dst is None: os.makedirs("plots", exist_ok=True) dst = os.path.join( - os.getcwd(), "plots", "{}_{}".format(int(time.time()), "timeline.png") + os.getcwd(), "plots", "{}_{}".format(int(time.time()), f"timeline.{format}") ) else: dst = os.path.expanduser(dst) if "~" in dst else dst - dst = "{}/{}".format(os.path.realpath(dst), "timeline.png") + dst = "{}/{}".format(os.path.realpath(dst), f"timeline.{format}") fig.savefig(dst)