def export_timeline_figs()

in smallpond/execution/scheduler.py [0:0]

• 73 lines of code
• 10 McCabe index (conditional complexity)

    def export_timeline_figs(self):
        from datetime import datetime

        import pandas as pd
        import plotly.express as px

        if self.large_runtime_state:
            logger.debug(f"pause exporting timeline figure")
            return

        now = datetime.now()
        task_data = pd.DataFrame(
            [
                dict(
                    task=repr(task),
                    node=(repr(node) if (node := self.logical_nodes.get(task.node_id, None)) is not None else "StandaloneTasks"),
                    status=str(task.status),
                    executor=task.exec_id,
                    start_time=datetime.fromtimestamp(task.start_time),
                    finish_time=datetime.fromtimestamp(
                        max(
                            task.finish_time or now.timestamp(),
                            task.start_time + 0.0001,
                        )
                    ),
                    elapsed_time=task.elapsed_time,
                    partition=str(task.partition_infos),
                    cpu_limit=task.cpu_limit,
                    gpu_limit=task.gpu_limit,
                    mem_limit=task.memory_limit,
                )
                for task in {**self.scheduled_tasks, **self.finished_tasks}.values()
                if task.start_time is not None
            ]
        )

        if task_data.empty:
            return

        timeline_figs = [
            px.timeline(
                task_data,
                x_start="start_time",
                x_end="finish_time",
                y="node",
                color="executor",
                hover_name="task",
                hover_data=task_data.columns,
                title="plan_timeline - progress: {}/{} ({:.1f}%), elapsed: {:.1f} secs, job: {}".format(
                    *self.progress, self.elapsed_time, self.ctx.job_id
                ),
                opacity=0.3,
            ),
            px.timeline(
                task_data,
                x_start="start_time",
                x_end="finish_time",
                y="executor",
                color="node",
                hover_name="task",
                hover_data=task_data.columns,
                title="exec_timeline - progress: {}/{} ({:.1f}%), elapsed: {:.1f} secs, job: {}".format(
                    *self.progress, self.elapsed_time, self.ctx.job_id
                ),
                opacity=0.3,
            ),
        ]

        for fig in timeline_figs:
            fig_title = str(fig.layout["title_text"])
            fig_filename, _ = fig_title.split(" - ", maxsplit=1)
            fig_filename += ".html"
            fig_path = os.path.join(self.ctx.log_root, fig_filename)
            fig.update_yaxes(autorange="reversed")  # otherwise tasks are listed from the bottom up
            fig.update_traces(marker_line_color="black", marker_line_width=1, opacity=1)
            fig.write_html(fig_path, include_plotlyjs="cdn" if pytest_running() else True)
            if self.ctx.shared_log_root:
                shutil.copy(fig_path, self.ctx.shared_log_root)
            logger.debug(f"exported timeline figure to {fig_path}")