Class TensorBoardParser

tests/tools.py:409–459

class TensorBoardParser:
    def __init__(self, log_dir: str):
        self.log_dir = log_dir
        self._event_files = self._find_event_files(log_dir)
        self._metrics = self._load_metrics()

    def _find_event_files(self, log_dir: str) -> List[str]:
        event_files = []
        for root, _, files in os.walk(log_dir):
            for f in files:
                if f.startswith("events.out.tfevents."):
                    event_files.append(os.path.join(root, f))
        return event_files

    def _load_metrics(self) -> Dict[str, Dict[int, float]]:
        metrics = defaultdict(dict)

        for event_file in self._event_files:
            ea = EventAccumulator(event_file)
            ea.Reload()
            tags = ea.Tags()["scalars"]
            for tag in tags:
                scalars = ea.Scalars(tag)
                for scalar in scalars:
                    step = scalar.step
                    value = scalar.value
                    if step not in metrics[tag] or value > metrics[tag][step]:
                        metrics[tag][step] = value
        return dict(metrics)

    def metric_exist(self, metric_name: str) -> bool:
        return metric_name in self._metrics

    def metric_min_step(self, metric_name: str) -> int:
        return min(self.metric_steps(metric_name))

    def metric_max_step(self, metric_name: str) -> int:
        return max(self.metric_steps(metric_name))

    def metric_steps(self, metric_name: str) -> List[int]:
        if not self.metric_exist(metric_name):
            raise ValueError(f"Metric '{metric_name}' does not exist.")
        return list(self._metrics[metric_name].keys())

    def metric_values(self, metric_name: str) -> List:
        if not self.metric_exist(metric_name):
            raise ValueError(f"Metric '{metric_name}' does not exist.")
        return list(self._metrics[metric_name].values())

    def metric_list(self, metric_prefix: str) -> List[str]:
        return [name for name in self._metrics if name.startswith(metric_prefix)]


class RayCleanupPlugin: