CspReporter.py

# Copyright (c) 2021 PaddlePaddle Authors. 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.

import os
import glob
import argparse

from multiprocessing import Process

from DCGMFileReader import dcgmFileReader
from ProfileFileReader import profileFileReader

from CspFileReader import getLogger
from CspFileReader import TIME_PATH, DCGM_PATH, NET_PATH, PROFILE_PATH
from CspFileReader import FILEORGANIZEFORM_BYRANK, FILEORGANIZEFORM_BYTRAINER


def get_argparse():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--profile_path',
                        type=str,
                        default='.',
                        help='Working path that store the monitor data.')

    parser.add_argument('--timeline_path',
                        type=str,
                        default='.',
                        help='Output timeline file name.')

    parser.add_argument('--gpuPerTrainer',
                        type=int,
                        default=8,
                        help='Gpus per trainer.')

    parser.add_argument('--trainerNum',
                        type=int,
                        default=4,
                        help='Num of trainer.')

    parser.add_argument('--groupSize',
                        type=int,
                        default=8,
                        help='Num of trainer in a group.')

    parser.add_argument('--displaySize',
                        type=int,
                        default=2,
                        help='Num of line need to display in a group.')

    return parser.parse_args()


class CspReporter(object):

    def __init__(self, args):
        self._args = args
        print(self._args)

        self._workPath = self._args.profile_path
        self._saveFilePath = self._args.timeline_path
        self._gpuPerTrainer = self._args.gpuPerTrainer
        self._groupSize = self._args.groupSize
        self._displaySize = self._args.displaySize
        self._trainerNum = self._args.trainerNum

        self._checkArgs()

        self._init_logger()
        self._init_timeInfo()
        self._init_reader()

    def _checkArgs(self):
        if self._trainerNum % self._groupSize != 0:
            raise Exception(
                "Input args error: trainerNum[%d] %% groupSize[%d] != 0" %
                (self._trainerNum, self._groupSize))

    def _init_logger(self):
        self._logger = getLogger()

    def _init_reader(self):
        self._dcgmPath = os.path.join(self._workPath, DCGM_PATH)
        self._netPath = os.path.join(self._workPath, NET_PATH)
        self._profilePath = os.path.join(self._workPath, PROFILE_PATH)

        self._netFileReaderArgs = {
            "dataPath": self._netPath,
            "groupSize": self._groupSize,
            "displaySize": self._displaySize,
            "gpuPerTrainer": self._gpuPerTrainer,
            "minTimeStamp": self._minTimeStamp,
            "organizeForm": FILEORGANIZEFORM_BYTRAINER,
        }

        self._dcgmFileReaderArgs = {
            "dataPath": self._dcgmPath,
            "groupSize": self._groupSize,
            "displaySize": self._displaySize,
            "gpuPerTrainer": self._gpuPerTrainer,
            "minTimeStamp": self._minTimeStamp,
            "organizeForm": FILEORGANIZEFORM_BYTRAINER,
        }

        self._profileFileReaderArgs = {
            "dataPath": self._profilePath,
            "groupSize": self._groupSize,
            "displaySize": self._displaySize,
            "gpuPerTrainer": self._gpuPerTrainer,
            "minTimeStamp": self._minTimeStamp,
            "organizeForm": FILEORGANIZEFORM_BYRANK,
        }

        self._dcgmFileReader = dcgmFileReader(self._logger,
                                              self._dcgmFileReaderArgs)
        self._profileFileReader = profileFileReader(self._logger,
                                                    self._profileFileReaderArgs)

    def _init_timeInfo(self):
        self._timePath = os.path.join(self._workPath, TIME_PATH)
        self._timeInfo = {}
        self._minTimeStamp = 0
        self._set_timeInfo()

    def _set_timeInfo(self, timeFileNamePrefix="time.txt", sed="."):
        timeFileNameList = glob.glob(
            os.path.join(self._timePath, timeFileNamePrefix, sed, "*"))
        for timeFileName in timeFileNameList:
            trainerId = int(timeFileName.split(sed)[-1])
            gpuId = int(timeFileName.split(sed)[-2])
            info = {}
            with open(timeFileName, "r") as rf:
                for line in rf:
                    if line.startswith("start time:"):
                        info["start_time"] = int(
                            float(line.split(":")[-1]) * 1e9)

                        self._minTimeStamp = min(self._minTimeStamp,
                                                 info["start_time"])

                    if line.startswith("end time:"):
                        info["end_time"] = int(float(line.split(":")[-1]) * 1e9)
            if not info:
                self._timeInfo[gpuId * trainerId] = info

    def _generateTraceFileByGroupAndGpuId(self, pipileInfo, netInfo, groupId,
                                          gpuId):
        dcgmInfoDict = self._dcgmFileReader.getDcgmInfoDict(groupId, gpuId)
        opInfoDict = self._profileFileReader.getOpInfoDict(groupId, gpuId)

        traceObj = {}
        traceObj["traceEvents"] = pipileInfo[str(
            gpuId)] + opInfoDict["traceEvents"] + dcgmInfoDict[
                "traceEvents"] + netInfo["traceEvents"]

        self._profileFileReader.dumpDict(traceObj, "traceFile", groupId, gpuId,
                                         False, self._saveFilePath)

    def _generateTraceFileByGroup(self, groupId, processNum):
        # first we need to generate pipeline info
        pipileInfo = self._profileFileReader.getPipeLineInfo(
            groupId, processNum)
        # second we need to generate dcgm info
        dcgmInfo = self._dcgmFileReader.getDCGMTraceInfo(groupId, processNum)

        # third we need to generate net info
        netInfo = {}
        netInfo["traceEvents"] = []
        # netInfo = self._netFileReader.parseFileByGroup(groupId, processNum)

        # forth we need to generate op info
        opInfo = self._profileFileReader.getOPTraceInfo(groupId)

        # finially we need dump this information into disk
        processPool = []
        pidList = []

        for gpuId in range(self._gpuPerTrainer):
            subproc = Process(target=self._generateTraceFileByGroupAndGpuId,
                              args=(
                                  pipileInfo,
                                  netInfo,
                                  groupId,
                                  gpuId,
                              ))
            processPool.append(subproc)
            subproc.start()
            pidList.append(subproc.pid)
            self._logger.info(
                "[traceFile]: process [%d] has been started, total task num is %d ..."
                % (subproc.pid, 1))

        for t in processPool:
            t.join()
            pidList.remove(t.pid)
            self._logger.info(
                "[traceFile]: process [%d] has exited! remained %d process!" %
                (t.pid, len(pidList)))

    def generateTraceFile(self, processNum=8):
        processPool = []
        pidList = []
        for groupId in range(self._trainerNum / self._groupSize):
            subproc = Process(target=self._generateTraceFileByGroup,
                              args=(
                                  groupId,
                                  processNum,
                              ))
            processPool.append(subproc)
            subproc.start()
            pidList.append(subproc.pid)
            self._logger.info(
                "[GroupTraceFile]: process [%d] has been started, total task num is %d ..."
                % (subproc.pid, 1))
        for t in processPool:
            t.join()
            pidList.remove(t.pid)
            self._logger.info(
                "[GroupTraceFile]: process [%d] has exited! remained %d process!"
                % (t.pid, len(pidList)))


if __name__ == '__main__':
    args = get_argparse()
    tl = CspReporter(args)
    tl.generateTraceFile()