Dubhe/dubhe-visual-server/parser_service/tbparser/graph_reader.py

"""
/**
* Copyright 2020 Zhejiang Lab. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================
*/
"""

# -*- coding: UTF-8 -*-
from tbparser.events_reader import EventReadingError, EventsFileReader
from tensorboard.compat.proto.graph_pb2 import GraphDef
from pathlib import Path
from typing import Union


class GraphReader:
    def __init__(self,
                 logdir: Union[str, Path],
                 stop_on_error: bool = False
    ):
        self._logdir = Path(logdir)
        self._stop_on_error = stop_on_error

    def read(self):
        log_files = sorted(f for f in self._logdir.glob('*') if f.is_file())
        for file_path in log_files:
            with open(file_path, 'rb') as f:
                reader = EventsFileReader(f)
                try:
                    result = self.read_graph_def(reader)
                    if result is not None:
                        return result
                except EventReadingError:
                    if self._stop_on_error:
                        raise
                    else:
                        continue
        return

    def read_graph_def(self, events):
        for event in events:
            if event.HasField('graph_def'):
                # wall_time = event.wall_time
                graph = GraphDef()
                graph.ParseFromString(event.graph_def)
                return graph
        return None